How to index a 3-D tensor with 2-D tensor in Tensorflow?

Question

I am trying to use a 2-D tensor to index a 3-D tensor in Tensorflow. For example, I have x of shape [2, 3, 4]:

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

 [[12, 13, 14, 15],
  [16, 17, 18, 19],
  [20, 21, 22, 23]]]

and I want to index it with another tensor y of shape [2, 3], where each element of y index the last dimension of x. For example, if we have y like:

[[0, 2, 3],
 [1, 0, 2]]

The output should of shape [2, 3]:

[[0, 6, 11],
 [13, 16, 22]]

Psidom Psidom · Accepted Answer · 2018-03-12T01:06:31

Create the indices using tf.meshgrid and then use tf.gather_nd to extract the elements:

# create a list of indices for except the last axis
idx_except_last = tf.meshgrid(*[tf.range(s) for s in x.shape[:-1]], indexing='ij')

# concatenate with last axis indices
idx = tf.stack(idx_except_last + [y], axis=-1)

# gather elements based on the indices
tf.gather_nd(x, idx).eval()

# array([[ 0,  6, 11],
#        [13, 16, 22]])

How to index a 3-D tensor with 2-D tensor in Tensorflow?

1 Answers