How to slice a tensor using given indices in tensorflow?

Question

I have a tensor with probabilities. This is a dynamic tensor with shape (?, 30) and I am selecting index with the best probability of these 30 values as :

    best_probability = tf.argmax(probability, axis = 1)

Now the dimensions of tensor best_probability is (?,). Now I want to select the values with these indices from another tensor called data with dimensions (?, 30, 1024, 3). Essentially from each of the 30 values select one with best probability using best_probability tensor.

The final output should have dimensions of (?, 1024, 3).

PS:- I tried gather_nd but it need indexing of best_probability tensor something like [[0, 9], [1, 10], [2, 15], [3, 25]]. To do so I wrote following snippet.

 selected_data = tf.stack(tf.range(probability.shape[0]),
                                tf.argmax(probability, axis = 1))

This doesn't work as I am dealing with a dynamic tensor. Is there any alternative to solve this problem.

Saurabh Pradhan Saurabh Pradhan · Accepted Answer · 2018-11-10T14:24:35

I was able to solve this issue using tf.batch_gather and tf.reshape

selected_data = tf.reshape(tf.batch_gather(data, best_probability),
                           (-1, data.shape[2],data.shape[3]))

How to slice a tensor using given indices in tensorflow?

1 Answers