python - 如何在纯 Tensorflow 中重写 Keras Stateful LSTM？

Question

任何人都可以分享关于如何将 Keras 有状态 LSTM 转换为纯 Tensorflow 模型的想法/博客/代码片段？然后批量训练它..

Tensorflow 不支持TPU 上的 Keras Stateful LSTM。他们的开发人员拒绝修复它。

我保留了大量的 TPU 时间，暂时无法使用它。任何帮助表示赞赏。

模型示例和训练代码：

model = Sequential()
model.add(LSTM(neurons, batch_input_shape=(window_size, n_steps, inputs_n), stateful=True)) 
model.add(Dense(outputs_n, activation='sigmoid'))
…
H = model.train_on_batch(X, y)

GitHub问题：https ://github.com/tensorflow/tensorflow/issues/28837