Selecting values from tensor based on an index tensor

user_6920649 0
Viewed 281

I have two matrices. Matrix A is contains some values and matrix B contains indices. The shape of matrix A and B is (batch, values) and (batch, indices), respectively.

My goal is to select values from matrix A based on indices of matrix B along the batch dimension.

For example:

# Matrix A
<tf.Tensor: shape=(2, 5), dtype=float32, numpy=
array([[0., 1., 2., 3., 4.],
       [5., 6., 7., 8., 9.]], dtype=float32)>

# Matrix B
<tf.Tensor: shape=(2, 2), dtype=int32, numpy=
array([[0, 1],
       [1, 2]], dtype=int32)>

# Expected Result
<tf.Tensor: shape=(2, 2), dtype=float32, numpy=
array([[0., 1.],
       [6., 7.]], dtype=int32)>

How can I achieve this in Tensorflow?

Many thanks in advance!

tensorflow2.0 tensorflow
1 Answers

You can achieve this with the tf.gather function.

mat_a = tf.constant([[0., 1., 2., 3., 4.],
                     [5., 6., 7., 8., 9.]])
mat_b = tf.constant([[0, 1], [1, 2]])

out = tf.gather(mat_a, mat_b, batch_dims=1)
out.numpy()
array([[0., 1.],
       [6., 7.]], dtype=float32)
user_92157800
Related