如何将tensorflow.Variable中的某些元素取出组成一个新的矩阵

在神经网络计算过程中，经常会遇到需要将矩阵中的某些元素取出并且单独进行计算的步骤（例如MLE，Attention等操作）。那么在 tensorflow 的 Variable 类型中如何做到这一点呢？

首先假设 Variable 是一个一维数组 A：

import numpy as np

import tensorflow as tf

a = np.array([1, 2, 3, 4, 5, 6, 7, 8])

A = tf.Variable(a)

我们把我们想取出的元素的索引存到 B 中，如果我们只想取出数组 A 中的某一个元素，则 B 的设定为：

b = np.array([3])

B = tf.placeholder(dtype=tf.int32, shape=[1])

由于我们的索引坐标只有一维，所以 shape=1。

取出元素然后组合成tensor C 的操作如下：

C = tf.gather_nd(A, B)

运行：

init = tf.global_variables_initializer()

with tf.Session() as sess:
    init.run()
    feed_dict = {B: b}
    result = sess.run([C], feed_dict=feed_dict)
    print result

得到：

[4]

如果我们想取出一维数组中的多个元素，则需要把每一个想取出的元素索引都单独放一行