如何基于TensorFlow大张量和指定列索引生成对应小张量？

你要实现沿张量第二维筛选指定索引的元素，直接给tf.gather指定axis=1参数即可，无需额外手动reshape，操作逻辑完全匹配需求：

• 筛选逻辑：自动取tensor_big第二维中所有在indices列表内的坐标对应的元素
• 输出形状：原生返回(3136, 410)的tensor_small，和预期形状完全一致

修正后的完整可运行代码如下：

import tensorflow as tf

if __name__ == '__main__':
    tensor_big = tf.random.uniform(
        (3136,512), minval=0, maxval=None, dtype=tf.dtypes.float32, seed=None, name=None
    )

    indices = tf.random.uniform(shape=[410, ], minval=0, maxval=512, dtype=tf.dtypes.int32, seed=None, name=None)
    # 核心修改：指定axis=1沿第二维取索引对应的元素
    tensor_small = tf.gather(tensor_big, indices, axis=1)

    # 输出验证形状：(3136, 410)
    print(tensor_small.shape)
    print(tensor_small)

如果你的TensorFlow版本支持，也可以用更直观的类Numpy索引写法实现相同效果：

tensor_small = tensor_big[:, indices]

两种写法运行结果完全一致，可根据代码习惯自行选择。

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