python - 純粋な Tensorflow で Keras Stateful LSTM を書き直す方法は?

Question

Keras Stateful 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