Tensor Entry Selection Logic Divergence in PyTorch & Numpy
Description

I'm setting up a torch.Tensor for masking purpose. When attempting to select entries by indices, it turns out that behaviors between using numpy.ndarray and torch.Tensor to hold index data are different. I would like to have access to the design in both frameworks and related documents that explain the difference.

Steps to replicate

Environment

Pytorch 1.3 in container from official release: pytorch/pytorch:1.3-cuda10.1-cudnn7-devel

Example

Say I need to set up mask as torch.Tensor object with shape [3,3,3] and set values at entries (0,0,1) &amp; (1,2,0) to 1. The code below explains the difference.

mask = torch.zeros([3,3,3])
indices = torch.tensor([[0, 1],
                        [0, 2],
                        [1, 0]])

mask[indices.numpy()] = 1 # Works
# mask[indices] = 1 # Incorrect result


I noticed that when using mask[indices.numpy()] a new torch.Tensor of shape [2], while mask[indices] returns a new torch.Tensor of shape [3, 2, 3, 3], which suggests difference in tensor slicing logic. 

You get different results because that's how indexing is implemented in Pytorch. If you pass an array as index, then it gets "unpacked". For example:

indices = torch.tensor([[0, 1], [0, 2], [1, 0]])

mask = torch.arange(1,28).reshape(3,3,3)

# tensor([[[ 1,  2,  3],
#          [ 4,  5,  6],
#          [ 7,  8,  9]],

#         [[10, 11, 12],
#          [13, 14, 15],
#          [16, 17, 18]],

#         [[19, 20, 21],
#          [22, 23, 24],
#          [25, 26, 27]]])


The mask[indices.numpy()] is equivalent to mask[[0, 1], [0, 2], [1, 0]], i.e. the elements of the i-th row of indices.numpy() are used to select elements of mask along i-th axis. So it returns tensor([mask[0,0,1], mask[1,2,0]]), i.e. tensor([2, 16]).

On the other hand, when passing a tensor as index (I don't know the exact reason for this differentiation between arrays and tensors for indexing), it is not "unpacked" like an array, and the elements of the i-th row of the indices tensor are used for selecting the elements of mask along the axis-0. That is, mask[indices] is equivalent to mask[[[0, 1], [0, 2], [1, 0]], :, :]

&gt;&gt;&gt; mask[ind]

tensor([[[[ 1,  2,  3],
          [ 4,  5,  6],
          [ 7,  8,  9]],

         [[10, 11, 12],
          [13, 14, 15],
          [16, 17, 18]]],


        [[[ 1,  2,  3],
          [ 4,  5,  6],
          [ 7,  8,  9]],

         [[19, 20, 21],
          [22, 23, 24],
          [25, 26, 27]]],


        [[[10, 11, 12],
          [13, 14, 15],
          [16, 17, 18]],

         [[ 1,  2,  3],
          [ 4,  5,  6],
          [ 7,  8,  9]]]])


which is basically tensor(mask[[0,1], :, :], mask[[0,2],: ,:], mask[[1,0], :, :]) and has shape indices.shape + mask[0,:,:].shape == (3,2,3,3). So whole "sheets" are selected and stacked into new dimensions. Note that this is not a new tensor, but a special view of mask. Therefore if you assign mask[indices] = 1, with this particular indices, then all the elements of mask will become 1.
