index: fix index of 0-element tensor by 0-element tensor.

[ysiraichi]

This PR adds support for indexing a 0-element tensor with a 0-element tensor index. It also adds a fast path whenever there are 0-element tensor indices. If found, we know that we will return a 0-element tensor.

In summary, here are the steps of this fast path:

• Assume there is, at least, one 0-element tensor index
• Deduce the output shape from the indices and start_dim
  • The output shape will be a concatenation of 3 sub-sequences:
    1. The shape of the indexed tensor up to dimension start_dim
    2. The shape of the indices
    3. The shape of the indexed tensor after dimension start_dim + len(indices)

For further details, see PyTorch's implementation of output shape computation.

cc @miladm @JackCaoG @lezcano

[lezcano]

We have a meta implementation of the meta of index_Tensor at https://github.com/pytorch/pytorch/blob/ee6cb6daa173896f8ea1876266a19775aaa4f610/torch/_meta_registrations.py#L2983

Perhaps it's better to replicate its behaviour here, as it's quite tricky.

I assume we cannot use this information here, even though dynamo knows about it, right?

[ysiraichi]

I think that I'm basically doing the same thing that this part of the meta implementation is doing. The only thing is that, at this point, we have already preprocessed the inputs such that:

• base is permuted so that the indexing dimensions are adjacent
• start_dim is the dimension that indices[0] indexes
• indices is a list of defined (i.e. tensor.defined() == true) tensors that indexes adjacent dimensions starting from start_dim

I assume we cannot use this information here, even though dynamo knows about it, right?

I think the problem is two-fold:

1. This function lives in Python. Thus, we need some way of calling it (e.g. this PoC)

2. This function only retrieves the output, throwing away all preprocess done in the beginning of it

   • For lowering, we would need to do the pre-processing, again

[lezcano]

Sounds reasonable. Thank you for looking into this!

[ysiraichi]

Thank you, Mario. I think I will wait for @JackCaoG review, too, before merging this PR.

[JackCaoG]

Let me take a look today