Как я могу получить доступ к состояниям из верхнего уровня при вычислении шлюзов в Tensorflow RNN?

Question

Я бы хотел изменить LSTM Tensorflow, чтобы учесть скрытое состояние верхнего слоя на предыдущем временном шаге при расчете значений гейта.

Я построил следующую модель:

inputs = tf.keras.Input(shape=(timesteps, feature_size))

l = customLSTM(units=state_size,
               return_sequences=True)(inputs)
l = customLSTM(units=state_size,
               return_sequences=True)(l)

output = TimeDistributed(Dense(output_size, activation='softmax'))(l)

model = tf.keras.Model(inputs=inputs, outputs=output)

"customLSTM" - это в основном слой "LSTM" Tensorflow, где я называю адаптированную версию "standard_lstm" 1, В этой пошаговой функции состояния h (t-1) и x (t) умножаются на их соответствующие веса и суммируются:

z = K.dot(x_t, kernel)
z += K.dot(h_tm1, recurrent_kernel)

Я хотел бы добавить другое взвешенное состояние h (l + 1, t-1), то есть скрытое состояние предыдущего временного шага и выше слоя:

z += K.dot(h_lp1_tm1, another_kernel)

Однако я не знаю, как я могу получить доступ к скрытому состоянию с верхнего уровня на предыдущем временном шагу. Каков наилучший способ сделать это?