tensorflow - 使用 Tensor 切片 Tensorflow Tensor

Question

我正在尝试使用在此 PR 中添加的“高级”、numpy 样式的切片，但是我遇到了 same issue as the user here :

ValueError: Shape must be rank 1 but is rank 2 for 'strided_slice_15' (op: 'StridedSlice') with input shapes: [3,2], [1,2], [1,2], [1].

即我想做这个 numpy 操作的等价物(在 numpy 中工作):

A = np.array([[1,2],[3,4],[5,6]]) 
id_rows = np.array([0,2])
A[id_rows]

但是，对于上述错误，这在 TF 中不起作用:

A = tf.constant([[1,2],[3,4],[5,6]])
id_rows = tf.constant([0,2])
A[id_rows]

Answer 1

你正在寻找这样的东西:

A = tf.constant([[1,2],[3,4],[5,6]])
id_rows = tf.constant([[0],[2]]) #Notice the brackets
out = tf.gather_nd(A,id_rows)