Say I want to index a tensor along an axis, but I don’t know a priori which axis. It could be:

`tensor[:, idx, :]`

Or it could be:

`tensor[idx, :, :]`

Or:

`tensor[:, :, idx]`

In these examples, the tensor has 3 possible axes. Is there a way to do this kind of indexing where both the axis and the index are variable?

You can use `tensor.index_select`

as seen here:

```
x = torch.randn(10, 10, 10)
idx = torch.randint(0, 10, (2,))
ref1 = x[idx]
ref2 = x[:, idx]
ref3 = x[:, :, idx]
out1 = x.index_select(0, idx)
out2 = x.index_select(1, idx)
out3 = x.index_select(2, idx)
print((ref1 == out1).all())
# tensor(True)
print((ref2 == out2).all())
# tensor(True)
print((ref3 == out3).all())
# tensor(True)
```