Tensorflow：如何使用另一个未知维度的张量索引一个张量？

Question

I have a tensor params with shape (?, 70, 64) and another tensor indices with shape (?, 1) . I want to index into the first tensor's axis 1 using the second tensor, to get a result with shape (?, 64) .

I can't figure how to go about it. Here's what I've tried:

tf.gather(params, indices)           # returns a tensor of shape (?, 1, 70, 64)
tf.gather(params, indices, axis=1)   # returns a tensor of shape (?, ?, 1, 64)
tf.gather_nd(params, indices)        # returns a tensor of shape (?, 70, 64)

(I have an older version of TensorFlow, which doesn't have batch_gather . ) Any help would be appreciated.

Thanks!

Answer 1

You can use tf.stack to convert your indices to a tensor of shape (?, 2) with the first number in the second dimension being the batch number. Then using this new indices with tf.gather_nd should give you what you want if I understand your goal correctly.

Since your indices is a tensor of shape (?, 1) , batch_gather would give you (?, 1, 64) , meaning one reshape step from your expected result tensor of shape (?, 64) . The following code shows two methods give you the same result:

import numpy as np
import tensorflow as tf

params = tf.constant(np.arange(3*70*64).reshape(3, 70, 64))
init_indices = tf.constant([[2], [1], [0]])
indices = tf.stack(
    [tf.range(init_indices.shape[0]), tf.reshape(init_indices, [-1])],
    axis=1
)
output = tf.gather_nd(params, indices)
batch_gather = tf.reshape(tf.batch_gather(params, init_indices),
                          [params.shape[0], -1])

with tf.Session() as sess:
    print('tf.gather_nd')
    print(output.shape)
    print(sess.run(output))
    print('batch_gather')
    print(batch_gather.shape)
    print(sess.run(batch_gather))

Edit on comment "first dimension unknown"

Overall, the optimal solution depends on the specific use case, and to use tf.gather_nd with tf.stack , the key is to get the batch size, ie the first dimension. One way, which again may not be optimal, is to use tf.shape :

import numpy as np
import tensorflow as tf

params = tf.placeholder(shape=(None, 70, 64), dtype=tf.int32)
init_indices = tf.placeholder(shape=(None, 1), dtype=tf.int32)
indices = tf.stack(
    [tf.range(tf.shape(init_indices)[0]), tf.reshape(init_indices, [-1])],
    axis=1
)
output = tf.gather_nd(params, indices)
batch_gather = tf.reshape(tf.batch_gather(params, init_indices),
                          [tf.shape(params)[0], -1])

with tf.Session() as sess:
    print('tf.gather_nd')
    print(output.shape)
    print(sess.run(
        output, feed_dict={params: np.arange(3*70*64).reshape(3, 70, 64),
                           init_indices: [[2], [1], [0]]}
    ))
    print('batch_gather')
    print(batch_gather.shape)
    print(sess.run(
        batch_gather, feed_dict={params: np.arange(3*70*64).reshape(3, 70, 64),
                                 init_indices: [[2], [1], [0]]}
    ))

One thing to point out is because batch size is unknown, print(batch_gather.shape) gives (?, ?) rather than (?, 64) .