Pytorch-index on multiple dimension tensor in a batch
Given a 1-d tensor:
A = torch.tensor([1, 2, 3, 4])

suppose we have some &quot;indexer tensor&quot;
ind1 = torch.tensor([3, 0, 1])
ind2 = torch.tensor([[3, 0], [1, 2]])

as we run A[ind1] &amp; A[ind2]
we get results tensor([4, 1, 2]) &amp; tensor([[4, 1],[2, 3]])
which is the same shape of the indexed tensor (ind1 and ind2) and its value are mapped from tensor A.
I want to ask how can I index on higher dimension tensors?

Currently I have one solution:
For a N-d tensor A, suppose we have the indexer tensor IND,
IND is like [[i11, i12, ... i1N], [i21, i22, ... i2N], ...[iM1, i22, ... iMN], where M is the number of indexed elements.
We can divide IND into N tensors, where
IND_1 = torch.tensor([i11, i21, ... iM1])
...
IND_N = torch.tensor([i1N, i2N, ... iMN])

as we run A[IND_1, ... IND_N], we got tensor(v1, v2, ... vM)
Example:
A = tensor([[1, 2], [3, 4]], [[5, 6], [7, 8]]])  # [2 * 2 * 2]
ind1 = tensor([1, 0, 1])
ind2 = tensor([1, 1, 0])
ind3 = tensor([0, 1, 0])
A[ind1, ind2, ind3]
=&gt; tensor([7, 4, 5])
# and the good thing is you can control the shape of result tensor by modifying the inds' shape.
ind1 = tensor([[0, 0], [1, 0]])
ind2 = tensor([[1, 1], [0, 1]])
ind3 = tensor([[0, 1], [0, 0]])
A[ind1, ind2, ind3]
=&gt; tensor([[3, 4],[5, 3]])  # same as inds' shape

Anyone has more elegant solutions?

1- Manual approach using unraveled indices on flattened input.
If you want to index on an arbitrary number of axes (all axes of A) then one straightforward approach is to flatten all dimensions and unravel the indices. Let's assume that A is 3D and we want to index it using a stack of ind1, ind2, and ind3:
&gt;&gt;&gt; ind = torch.stack((ind1, ind2, ind3))

You can first unravel the indices using A's strides:
&gt;&gt;&gt; unraveled = torch.tensor(A.stride()) @ ind.flatten(1)

Then flatten A, index it with unraveled and reshape to the final form:
&gt;&gt;&gt; A.flatten()[unraveled].reshape_as(ind[0])


2- Using a simple split of ind.
You can actually perform the same operation using torch.chunk:
&gt;&gt;&gt; A[ind.chunk(len(ind))][0]

Or alternatively torch.split which is identical:
&gt;&gt;&gt; A[ind.split(1)][0]


3- Initial answer for single-axis indexing.
Let's take a minimal multi-dimensional example with A being a 2-D tensor defined as:
&gt;&gt;&gt; A = torch.tensor([[1, 2, 3, 4],
                      [5, 6, 7, 8]])

From your description of the problem:

the same shape of index tensor and its value are mapped from tensor A.

Then the indexer tensor would require to have the same shape as the indexed tensor A, since this one is no longer flat. Otherwise, what would the result of A (shaped (2, 4)) indexed by ind1 (shape (3,)) be?
If you are indexing on a single dimension then you can utilize torch.gather:
&gt;&gt;&gt; A.gather(1, ind2)
tensor([[4, 1],
        [6, 7]])


