[feature request] "Batched" index_select (i.e. simplified torch.gather with not specifying full index)

[vadimkantorov]

E.g. we have I = torch.randint(0, n3, (n1, n2)) and T = torch.rand(n1, n2, n3, n4, n5)
We'd like to compute O[i, j, ...] = T[i, j, I[i, j], ...]

This is fairly frequent. Currently this is possible with a combination of multi unsqueeze + expand + gather + squeeze.

I propose a new helper function to do this in one go (and thus making this pattern less complex and error-prone) or just making gather support this directly by relaxing index tensor shape constraints

cc @ezyang @gchanan @zou3519 @mikaylagawarecki

[vadimkantorov]

Real-world example: https://github.com/pytorch/vision/blob/7947fc8fb38b1d3a2aca03f22a2e6a3caa63f2a0/torchvision/models/detection/roi_heads.py#L129

Here it is for 1D index, then this functionality is possible by using advanced indexing with two tensors. But for larger dimensions this is not possible. Also, this particular form of advanced indexing is not very well-known.

[vadimkantorov]

My implementation:

gather_incomplete_left = lambda tensor, I: tensor.gather(I.ndim, I[(...,) + (None,) * (tensor.ndim - I.ndim)].expand((-1,) * (I.ndim + 1) + tensor.shape[I.ndim + 1:])).squeeze(I.ndim)

In addition to being quite convoluted, this would also quite likely fail in JIT

Related issues: #7991, #62352

[vadimkantorov]

This could also be used for generalizing F.cross_entropy accepted batch shapes

[vadimkantorov]

Alternatively, index_select could be extended somehow to allow >1dim index, maybe naming could be index_left

[zou3519]

This is an interesting ask. Are there other libraries that have something like this?

[vadimkantorov]

Never saw anything like this in libraries, but in practice this generalization of index_select / simplification of tather is very frequent

[vadimkantorov]

So one could think of this as batched index_select

[vadimkantorov]

cc @Chillee