Инициализация состояния RNN с использованием набора данных и / или оценки тензорного потока

Question

Я хотел бы использовать tf.data и tf.estimator как часть решения для обучения сети RNN (например, LSTM).

RNN будет использоваться для прогнозирования продолжения последовательности.Длинная входная последовательность X будет разбита на отдельные примеры X (i).

Входом в сеть будет X (i), а скрытое состояние будет выводиться сетью после подачи X (i-1).Выходными данными будут Y (i) и новое скрытое состояние S (i):

         Y(i)
           | 
S(i-1) -> RNN -> S(i)
           |
         X(i)

Я хочу поэкспериментировать со следующей настройкой:

                     Y(i)
S(i-100)               | 
S(i-50)  -> E(i)-> -> RNN -> S(i)
S(i-1)                 |
                     X(i)

В каком S (i-100), S (i-50), S (i-1) - это выходы скрытого состояния с шагом соответственно 100, 50 и 1 раз.E (i) - это изученное вложение.

Теперь возникает вопрос: как я могу использовать API tf.data и / или tf.estimator для получения скрытого предыдущего состояния S (ix) при обучении по X (я)?Предлагаемое решение должно позволить распределенное обучение.

Можно ли использовать переменную tf.Variable и назначить выходное состояние S (i) на каждом этапе обучения?Доступно ли значение переменной для всех экземпляров обучения во время распределенного обучения?