index.py

import torch

t = torch.tensor([
[1, 2, 3, 4, 5], 
[6,7,8,9,10],
[11,12,13,14,15],
[16,17,18,19,20]
]) # 4*5
i = torch.tensor([
[1,2,3], 
[3,2,1],
]) # 2*3

o = t[i] 
  = torch.tensor([
  [[ 6,  7,  8,  9, 10], [11, 12, 13, 14, 15], [16, 17, 18, 19, 20]],
  [[16, 17, 18, 19, 20], [11, 12, 13, 14, 15], [ 6,  7,  8,  9, 10]]
  ]) #  2*3*5
  

########## same function by tensorflow
import tensorflow as tf

t = tf.constant([
[1, 2, 3, 4, 5], 
[6,7,8,9,10],
[11,12,13,14,15],
[16,17,18,19,20]
]) # 4*5

i =  tf.constant([
[1,2,3], 
[3,2,1],
]) # 2*3

i_expand = tf.expand_dims(i,axis=2)
<tf.Tensor: shape=(2, 3, 1), dtype=int32, numpy=
array([[[1],
        [2],
        [3]],

       [[3],
        [2],
        [1]]], dtype=int32)>

io=tf.gather_nd(t,tf.expand_dims(i,axis=2))
<tf.Tensor: shape=(2, 3, 5), dtype=int32, numpy=
array([[[ 6,  7,  8,  9, 10],
        [11, 12, 13, 14, 15],
        [16, 17, 18, 19, 20]],

       [[16, 17, 18, 19, 20],
        [11, 12, 13, 14, 15],
        [ 6,  7,  8,  9, 10]]], dtype=int32)>

Multi Indexes with different shape

###############################################################################3
>>> t = torch.arange(5*4*3).view(5,4,3)
>>> t
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]],

        [[36, 37, 38],
         [39, 40, 41],
         [42, 43, 44],
         [45, 46, 47]],

        [[48, 49, 50],
         [51, 52, 53],
         [54, 55, 56],
         [57, 58, 59]]])
>>> t.shape
torch.Size([5, 4, 3])
>>> i1 = torch.randint(0,3,(3,3))
>>> i1
tensor([[1, 2, 0],
        [0, 2, 0],
        [0, 0, 2]])
>>> i2 = torch.randint(0,3,(3,))
>>> i2
tensor([1, 1, 0])
>>> o = t[i1,i2]
>>> o
tensor([[[15, 16, 17],
         [27, 28, 29],
         [ 0,  1,  2]],

        [[ 3,  4,  5],
         [27, 28, 29],
         [ 0,  1,  2]],

        [[ 3,  4,  5],
         [ 3,  4,  5],
         [24, 25, 26]]])
tensor([[t[1,1], t[2,1], t[0,0]],

        [t[0,1], t[2,1], t[0,0]],

        [t[0,1], t[0,1], t[2,0]]])
>>> o.shape
torch.Size([3, 3, 3])