tf.gather_nd 详解及例子

官网链接
https://tensorflow.google.cn/api_docs/python/tf/gather_nd

tf.gather_nd(
    params,
    indices,
    name=None
)

按照indices的格式从params中抽取切片（合并为一个Tensor）
indices是一个K维整数Tensor。

例子：

import tensorflow as tf

a = tf.Variable([[1, 2, 3, 4, 5],
                 [6, 7, 8, 9, 10],
                 [11, 12, 13, 14, 15]])
index_a1 = tf.Variable([[0, 2], [0, 4], [2, 2]])  # 随便选几个
index_a2 = tf.Variable([0, 1])  # 0行1列的元素——2
index_a3 = tf.Variable([[0], [1]])  # [第0行,第1行]
index_a4 = tf.Variable([0])  # 第0行


with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    print(sess.run(tf.gather_nd(a, index_a1)))
    print(sess.run(tf.gather_nd(a, index_a2)))
    print(sess.run(tf.gather_nd(a, index_a3)))
    print(sess.run(tf.gather_nd(a, index_a4)))

输出：

[ 3  5 13]
2
[[ 1  2  3  4  5]
 [ 6  7  8  9 10]]
[1 2 3 4 5]