tensorflow - 如何在 TensorFlow 中选择二维张量的某些列？

Question

由于正在 this issue 中进行广义切片，实现 2D 张量(矩阵)的 op 收集列的最佳方法是什么？例如，对于张量 t :

1 2 3 4
5 6 7 8 

和指数 [1,3]，我想得到:

2 4
6 8

相当于 numpy t[:, [1,3]] .

Answer 1

同时gather方法有一个 axis范围。

import tensorflow as tf
params = tf.constant([[1,2,3],[4,5,6]])
indices = [0,2]
op = tf.gather(params, indices, axis=1)

产生输出

[[1 3]
 [4 6]]