Torch gather middle dimension

user_8737016 0
Viewed 921

Let a be a (n, d, l) tensor. Let indices be a (n, 1) tensor, containing indices. I want to gather from a in the middle dimension tensors from indices given by indices. The resulting tensor would therefore be of shape (n, l).

n = 3
d = 2
l = 3

a = tensor([[[ 0,  1,  2],
             [ 3,  4,  5]],

            [[ 6,  7,  8],
             [ 9, 10, 11]],

            [[12, 13, 14],
             [15, 16, 17]]])

indices = tensor([[0],
                  [1],
                  [0]])

# Shape of result is (n, l)
result = tensor([[ 0,  1,  2],  # a[0, 0, :] since indices[0] == 0

                 [ 9, 10, 11],  # a[1, 1, :] since indices[1] == 1

                 [12, 13, 14]]) # a[2, 0, :] since indices[2] == 0

This is indeed similar to a.gather(1, indices), but gather won't work since indices does not have the same shape as a. How can I use gather in this setting? Or what should I use?

torch pytorch python
3 Answers

You can create the indices manually. The indices tensor has to be flattened if it has the shape of your example data.

a[torch.arange(len(a)),indices.view(-1)]
# equal to a[[0,1,2],[0,1,0]]

Out:

tensor([[ 0,  1,  2],
        [ 9, 10, 11],
        [12, 13, 14]])
user_142777220

before use gather function, reshape indices, this is a example

def gather_righthand(src, index, check=True):
    index = index.long()
    i_dim = index.dim(); s_dim = src.dim(); t_dim = i_dim-1
    if check:
        assert s_dim > i_dim
        for d in range(0, t_dim): 
            assert src.shape[d] == index.shape[d]
    index_new_shape = list(src.shape)
    index_new_shape[t_dim] = index.shape[t_dim]
    for _ in range(i_dim, s_dim): index = index.unsqueeze(-1)

    index_expand = index.expand(index_new_shape)            # only this two line matters
    return torch.gather(src, dim=t_dim, index=index_expand) # only this two line matters


gather_righthand(a, indices)
tensor([[[ 0.,  1.,  2.]],
        [[ 9., 10., 11.]],
        [[12., 13., 14.]]])
user_163516390

I add my answer on top of Michael's for more dimensions on either side of the indexed dimension, but I'd like someone to give me a better one that doesn't use arange!

def squeeze_index(x, dim, index):
  # flatten to rows
  y = x.view((-1,) + x.shape[dim:])

  # generate row indices
  rows = torch.arange(y.shape[0])

  # index and reshape
  result_shape = x.shape[:dim] + (x.shape[dim+1:] if dim != -1 else ())
  return y[rows, index.view(-1), ...].view(result_shape)

a = torch.arange(2*3*2*3).reshape((2,3,2,3))
indices = torch.tensor([0,0,1,0,0,1]).reshape((2,3))
result = squeeze_index(a, 2, i)
print("a", a.shape, a)
print("indices", indices.shape, indices)
print("result", result.shape, result)

Gives:

a torch.Size([2, 3, 2, 3]) tensor([[[[ 0,  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, 28, 29]],

         [[30, 31, 32],
          [33, 34, 35]]]])
indices torch.Size([2, 3]) tensor([[0, 0, 1],
        [0, 0, 1]])
result torch.Size([2, 3, 3]) tensor([[[ 0,  1,  2],
         [ 6,  7,  8],
         [15, 16, 17]],

        [[18, 19, 20],
         [24, 25, 26],
         [33, 34, 35]]])
user_4048610
Related