如何利用索引数组将(128,10,51)的TensorFlow张量转为(128,51)？

问题背景 #

现有形状为(128, 10, 51)的TensorFlow张量A，以及形状为(128,)、元素取值为0-9的ndarray索引数组B。需求是从张量A的第二维度中选取对应索引的元素，转换为形状为(128, 51)的张量，替代现有的for循环实现。

简洁实现方案 #

单行代码方案 #

利用TensorFlow的tf.gather函数，结合batch_dims参数可直接完成批量索引选取：

result = tf.gather(A, tf.convert_to_tensor(B), axis=1, batch_dims=1)

代码说明 #

• tf.convert_to_tensor(B)：将ndarray类型的索引数组B转为TensorFlow张量，适配API要求。
• axis=1：指定要选取元素的目标维度（即张量A的第二维度，索引从0开始）。
• batch_dims=1：将前1个维度（(128,)的批次维度）设为批量维度，让每个批次样本对应选取B中匹配的索引元素，最终直接输出(128, 51)形状的张量。

内容的提问来源于stack exchange，提问作者Behnam Mahdinia