如何从二维tensorflow数组中检索特定元素?

Question

我有一个形状为(4,2)的张量

<tf.Tensor: shape=(4, 2), dtype=int32, numpy=
array([[0, 1],
[2, 3],
[4, 5],
[6, 7]])>

如果我只能访问张量中的选择元素，比如1,2,5,6，我该怎么做呢，假设我有两个列表或numpy数组M =[0,1,2,3] AND N=[1,0,1,0]来指定元素我该如何访问它们，如果是numpy数组数组[M,N]就可以了但在tensorflow中它说:

Only integers, slices (`:`), ellipsis (`...`), tf.newaxis (`None`) and scalar tf.int32/tf.int64 tensors are valid indices, got array 

是否有特定的方法来实现!?

Answer 1

使用gather_nd():

a = tf.constant([[0, 1],
[2, 3],
[4, 5],
[6, 7]])
M = [0,1,2,3]
N = [1,0,1,0]
inds = tf.stack((M, N))
inds = tf.transpose(inds)
res = tf.gather_nd(a, inds)