我有一个Keras模型，该模型使用带有return_sequences = True的LSTM层，从而产生了形状为X的张量(None, sequence_len, feature_dim)。
使用另一个Y形状的张量(None, num_slices)，其中包含num_slices个索引值，我想在{的num_slices中的特定索引处动态选择Y个切片{1}}沿其第二轴（以X的方式）-产生形状为X[:, idx, :]的张量。
我正在寻找的功能的说明性示例如下：
如果(None, num_slices, feature_dim)对应形状为X的{{1}}层的张量输出，并且LSTM()是该(None, 3, 2)层的张量输出我们想要用形状Y切成第二个轴的索引对。
Input()我寻求的Keras层将LSTM和输入层张量作为参数，并产生(None, 2)的输出张量，该输出张量沿第二轴被X = tf.Variable([[[3,5],
[2,4],
[7,1]],
[[6,10],
[4,8],
[14,2]],
[[9,15],
[6,12],
[21,3]],
[[12,20],
[8,16],
[28,4]]])
Y = tf.Variable([[1, 2], [0,1], [1,1], [0,0]])
中的索引切开：
X也就是说，与Y Keras张量相对应。
我尝试使用<tf.Tensor: id=2901, shape=(4, 2, 2), dtype=int32, numpy=
array([[[2, 4],
[7, 1]],
[[ 6, 10],
[ 4, 8]],
[[ 6, 12],
[ 6, 12]],
[[12, 20],
[12, 20]]], dtype=int32)>
，但对我来说效果不佳。
我还尝试了(None, 2, 2)，它产生的形状是tf.gather_nd()而不是tf.gather(X, Y, axis = 1)形状。
但是，我确实注意到，(4, 4, 2, 2)的结果张量具有我在每个(4, 2, 2)切片中寻找的输出，但是无法将张量减小为所需的形状。
是否可以在自定义Keras层（或Lambda层）甚至原始的Tensorflow代码中获得预期的行为？
谢谢。