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

Question

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

Tensorflow在TPU上不支持Keras 有状态 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