Tensorflow / Keras: Dynami c индексирование трехмерного тензора с использованием тензора индексов

Question

У меня есть модель Keras, использующая слои LSTM с return_sequences = True, таким образом, получая тензор X формы (None, sequence_len, feature_dim).

Используя другой тензор Y формы (None, num_slices), содержащий num_slices значения индекса, я хочу динамически выбирать num_slices срезов по указанным c индексам в Y из X вдоль его второй оси (X[:, idx, :] образом) - получая тензор формы (None, num_slices, feature_dim).

Иллюстративный пример нужной мне функциональности следующий:

Если X соответствует тензорному выходу слоя LSTM() с формой (None, 3, 2) и Y - тензорный выход слоя Input() пар индекса, по которому мы хотим нарезать вторую ось с помощью формы (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]])

Слой Keras, который я ищу, принимает в качестве аргумента LSTM и Input Слои тензоров и получим выходной тензор X, нарезанный вдоль второй оси по индексам в Y:

<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).

Однако я заметил, что полученный тензор tf.gather(X, Y, axis = 1) имел выход, который я ищу в каждом срезе tf.gather(X, Y, axis = 1)[idx, idx, : , :], но не смог уменьшить тензор до требуемой формы.

Возможно ли получить ожидаемое поведение в пользовательском слое Keras (или в лямбда-слое) или даже в необработанном коде Tensorflow?

Заранее спасибо.