你可以使用tf.gather_nd得到你想要的结果。最后的表达式是：

tf.gather_nd(params, tf.stack([tf.tile(tf.expand_dims(tf.range(tf.shape(indices)[0]), 1), [1, tf.shape(indices)[1]]), tf.transpose(tf.tile(tf.expand_dims(tf.range(tf.shape(indices)[1]), 1), [1, tf.shape(indices)[0]])), indices], 2))

此表达式有以下解释：

• tf.gather_nd执行您所期望的操作，并使用索引从params收集输出
• tf.stack组合了三个独立的Tensor，最后一个Tensor是索引。前两个Tensor指定前两个维（参数/索引的轴0和轴1）的顺序。
• 对于所提供的示例，轴0的顺序简单地为0、1、2、...、63，轴1的顺序简单地为0、1、2、...、783。这些序列分别由tf.range(tf.shape(indices)[0])和tf.range(tf.shape(indices)[1])获得。
• 对于所提供的示例，索引具有形状（64，784）。来自上面最后一点的其他两个Tensor需要具有相同的形状，以便与tf.stack组合
• 首先，使用tf.expand_dims向两个序列中的每一个添加额外的维度/轴。
• tf.tile和tf.transpose的使用可通过示例显示：假设params和index的前两个轴具有形状（5，3）。我们希望第一Tensor为：

[[0, 0, 0], [1, 1, 1], [2, 2, 2], [3, 3, 3], [4, 4, 4]]

我们希望第二Tensor为：

[[0, 1, 2], [0, 1, 2], [0, 1, 2], [0, 1, 2], [0, 1, 2]]

这两个Tensor的功能几乎类似于为相关索引指定网格中的坐标。

• tf.stack的最后一部分将三个Tensor组合在一个新的第三个轴上，因此结果具有与params相同的3个轴。

请记住，如果您的轴比问题中的轴多或少，您需要相应地修改tf.stack中指定坐标的Tensor的数量。

你需要的是一个自定义的约简函数。如果你要保持indices索引，那么我建议使用tf.reduce_max：

max_params = tf.reduce_max(params_tensor, reduction_indices=[2])

否则，这里有一种方法可以得到你想要的（Tensor对象是不可赋值的，所以我们创建一个二维Tensor列表，并使用tf.pack打包它）：

import tensorflow as tf
import numpy as np
with tf.Graph().as_default():

    params_tensor = tf.pack(np.random.randint(1,256, [5,5,10]).astype(np.int32))

    indices = tf.pack(np.random.randint(1,10,[5,5]).astype(np.int32))

    output = [ [None for j in range(params_tensor.get_shape()[1])] for i in range(params_tensor.get_shape()[0])]  
    for i in range(params_tensor.get_shape()[0]):
        for j in range(params_tensor.get_shape()[1]):
            output[i][j] = params_tensor[i,j,indices[i,j]]
    output = tf.pack(output)

    with tf.Session() as sess:
        params_tensor,indices,output = sess.run([params_tensor,indices,output])

        print params_tensor
        print indices
        print output

我知道我迟到了，但是我最近不得不做一些类似的事情，并且能够使用Ragged Tensors来完成它：

output =  tf.gather(params, tf.RaggedTensor.from_tensor(indices), batch_dims=-1, axis=-1)

希望能有所帮助