Python 3.x 我怎样才能得到一个张量的一部分，其索引来自另一个张量？（Keras，tf后端）

我正在使用RNN和注意机制

不同于软注意（通过注意张量对RNN的时间步进行加权并求和），我想选择最重要的

k

时间步，并使用其他方法将它们组合起来

因此，我有三个张量：

• RNNMatrix

  ，具有形状

  （无，时间步长，单元格数）

  ，其中

  None

  是批次大小

• AttMatrix

  ，带形状

  （无，时间步长）

• IdxMatrix

  ，形状为

  （无，k）

  ，每行是最重要的k个时间步的索引，从

  AttMatrix

理想情况下，我可以通过

RNNMatrix

和

IdxMatrix

获得

KRNNMatrix

形状

（无，k，单元格数）

获取

IdxMatrix

的简单方法是使用

tf.nn.top\k

。我的代码如下：

selected_att_num = 10

def My_select_layer(ip, k_num):

    LSTM_out = ip[0]
    Attention = ip[1]

    idx = tf.nn.top_k(Attention, k=k_num, sorted = True).indices

    # gen row index
    row_idx = tf.range(tf.shape(Attention)[0])

    # repeat k times
    row_idx_repeat = tf.keras.backend.repeat(tf.reshape(row_idx,(-1,1)), k_num)

    # gen index for gather_nd
    row_idx_repeat_reshape = tf.reshape(row_idx_repeat, (-1,k_num))
    idx_stacked = tf.stack([row_idx_repeat_reshape, idx], axis=-1)

    return tf.gather_nd(LSTM_out, idx_stacked)


_input = Input(shape = (data_length, data_dim))
LSTM_out = LSTM(cells, return_sequences=True)(_input )

# compute att for each step
Attention = Activation('tanh')(LSTM_out)
Attention = Dense(1, use_bias=False)(Attention)
Attention = Flatten()(Attention)
Attention = Activation('softmax', name='attention')(Attention)

LSTM_out_gathered = Lambda(My_select_layer,arguments={'k_num':selected_att_num})([LSTM_out , Attention])

模型符合得很好，在我对其进行培训时出现了一个错误：

ValueError：一个操作对渐变有

None

。请确保所有操作都定义了梯度（即可微）。无梯度的普通操作：K.argmax、K.round、K.eval

我很确定这个错误归咎于我的Lambda层，很可能归咎于

tf.nn.top_k

的变量

index

但我不知道如何修复它

或者，“按索引选择张量”操作在理论上是不可微的吗？如果是，是否有解决办法？
在最坏的情况下，强化学习能保持这种状态吗

我也在网上搜索了很多，找到了丹尼尔·莫勒（Daniel Möller）提出的一个非常类似的问题。但是，Daniel没有详细说明他是如何解决这个问题的