Pytorch tensor indexing error for sizes M

I am trying to access a pytorch tensor by a matrix of indices and I recently found this bit of code that I cannot find the reason why it is not working.

The code below is split into two parts. The first half proves to work, whilst the second trips an error. I fail to see the reason why. Could someone shed some light on this?

import torch
import numpy as np

a = torch.rand(32, 16)
m, n = a.shape
xx, yy = np.meshgrid(np.arange(m), np.arange(m))
result = a[xx]   # WORKS for a torch.tensor of size M >= 32. It doesn't work otherwise.

a = torch.rand(16, 16)
m, n = a.shape
xx, yy = np.meshgrid(np.arange(m), np.arange(m))
result = a[xx]   # IndexError: too many indices for tensor of dimension 2

and if I change a = np.random.rand(16, 16) it does work as well.

CodePudding user response：

First, let me give you a quick insight into the idea of indexing a tensor with a numpy array and another tensor.

Example: this is our target tensor to be indexed

    numpy_indices = torch.tensor([[0, 1, 2, 7],
                                  [0, 1, 2, 3]])   # numpy array

    tensor_indices = torch.tensor([[0, 1, 2, 7],
                                   [0, 1, 2, 3]])   # 2D tensor

    t = torch.tensor([[1,  2,  3,   4],            # targeted tensor
                      [5,  6,  7,   8],
                      [9,  10, 11, 12],
                      [13, 14, 15, 16],
                      [17, 18, 19, 20],
                      [21, 22, 23, 24],
                      [25, 26, 27, 28],
                      [29, 30, 31, 32]])
     numpy_result = t[numpy_indices]
     tensor_result = t[tensor_indices]

Indexing using a 2D numpy array: the index is read like pairs (x,y) tensor[row,column] e.g. t[0,0], t[1,1], t[2,2], and t[7,3].
```
print(numpy_result)  # tensor([ 1,  6, 11, 32])
```

Indexing using a 2D tensor: walks through the index tensor in a row-wise manner and each value is an index of a row in the targeted tensor. e.g. [ [t[0],t[1],t[2],[7]] , [[0],[1],[2],[3]] ] see the example below, the new shape of tensor_result after indexing is (tensor_indices.shape[0],tensor_indices.shape[1],t.shape[1])=(2,4,4).

print(tensor_result)     # tensor([[[ 1,  2,  3,  4],
                         #          [ 5,  6,  7,  8],
                         #          [ 9, 10, 11, 12],
                         #          [29, 30, 31, 32]],

                         #          [[ 1,  2,  3,  4],
                         #           [ 5,  6,  7,  8],
                         #           [ 9, 10, 11, 12],
                         #           [ 13, 14, 15, 16]]])

If you try to add a third row in numpy_indices, you will get the same error you have because the index will be represented by 3D e.g., (0,0,0)...(7,3,3).

indices = np.array([[0, 1, 2, 7],
                    [0, 1, 2, 3],
                    [0, 1, 2, 3]])

print(numpy_result)   # IndexError: too many indices for tensor of dimension 2

However, this is not the case with indexing by tensor and the shape will be bigger (3,4,4).

Finally, as you see the outputs of the two types of indexing are completely different. To solve your problem, you can use

xx = torch.tensor(xx).long()  # convert a numpy array to a tensor

What happens in the case of advanced indexing (rows of numpy_indices > 3 ) as your situation is still ambiguous and unsolved and you can check 1 , 2, 3.

CodePudding user response：

To whoever comes looking for an answer: it looks like its a bug in pyTorch.

Indexing using numpy arrays is not well defined, and it works only if tensors are indexed using tensors. So, in my example code, this works flawlessly:

a = torch.rand(M, N)
m, n = a.shape
xx, yy = torch.meshgrid(torch.arange(m), torch.arange(m), indexing='xy')
result = a[xx]   # WORKS

I made a gist to check it, and it's available here