python - 在 TensorFlow 中索引子张量

Question

我想通过将index应用于输入张量的子张量来获得输出张量。

例如在 NumPy 中，

import numpy as np

input = np.random.random((100,5)) # matrix
index = np.randint(5, size=(100,)) # vector
output = data[np.arange(index.shape[0]), index] # vector

给我想要的输出(我想要这个的符号版本)。

与 Theano 类似，

import theano.tensor as T
import theano

input = T.matrix() # symbolic matrix
index = T.ivector() # symbolic vector
output = input[T.arange(index.shape[0]), index] # symbolic vector

给了我想要的输出。

如何在 TensorFlow 中执行此操作？

import tensorflow as tf
input = tf.placeholder('float32', [None, 5])
index = tf.placeholder('int32', [None])
output = ???

与 NumPy 的示例不同，index 的长度(=input 的第一个维度)不固定。

Answer 1

您可以使用 tf.gather_nd 进行切片:

output = tf.gather_nd(input, tf.stack((tf.range(tf.shape(index)[0]), index), -1))