How to use indices from tf.nn.top_k with tf.gather_nd?

user_4850133 0
Viewed 2310

I'm trying to use the indices returned from tf.nn.top_k to extract the values from a second tensor.

I've tried using numpy type indexing, as well as tf.gather_nd directly, but I noticed that the indexing is wrong.

#  temp_attention_weights of shape [I, B, 1, J]
top_values, top_indices = tf.nn.top_k(temp_attention_weights, k=top_k)

# top_indices of shape [I, B, 1, top_k], base_encoder_transformed of shape [I, B, 1, J]

# I now want to extract from base_encoder_transformed top_indices
base_encoder_transformed = tf.gather_nd(base_encoder_transformed, indices=top_indices)  

# base_encoder_transformed should be of shape [I, B, 1, top_k]

I noticed that top_indices is of the wrong format, but I can't seem to transform it to be used in tf.gather_nd, where the innermost dimension is used to index each corresponding element from base_encoder_transformed. Does anybody know a way to get top_indices into the right format?

tensorflow python
2 Answers

top_indices will index only over the last axis, you need to add indices for the rest of axes too. That is easy with tf.meshgrid:

import tensorflow as tf

# Example input data
I = 4
B = 3
J = 5
top_k = 2
x = tf.reshape(tf.range(I * B * J), (I, B, 1, J)) % 7
# Top K
top_values, top_indices = tf.nn.top_k(x, k=top_k)
# Make indices for the rest of axes
ii, jj, kk, _ = tf.meshgrid(
    tf.range(I),
    tf.range(B),
    tf.range(1),
    tf.range(top_k),
    indexing='ij')
# Stack complete index
index = tf.stack([ii, jj, kk, top_indices], axis=-1)
# Get the same values again
top_values_2 = tf.gather_nd(x, index)
# Test
with tf.Session() as sess:
    v1, v2 = sess.run([top_values, top_values_2])
    print((v1 == v2).all())
    # True
user_17827920

I don't see a reason to use tf.gather_nd. There is an easier and faster (no need to use tf.meshgrid) solution using tf.gather with the batch_dims parameter.

import tensorflow as tf

# Example input data
I = 4
B = 3
J = 5
top_k = 2
x = tf.reshape(tf.range(I * B * J), (I, B, 1, J)) % 7
# Top K
top_values, top_indices = tf.nn.top_k(x, k=top_k)
#Gather indices along last axis
top_values_2 = tf.gather(x, top_indices, batch_dims = 3)

tf.reduce_all(top_values_2 == top_values).numpy()
#True

Note that batch_dims is 3 in this case because we want to gather from the last axis, and the rank of x is 4.

user_117815460
Related