Извлечение состояний клеток для каждого временного шага в тензорном потоке

Question

Я работаю с данными последовательности переменной длины.Я пытался реализовать архитектуру с двумя LSTM.Предполагается, что второй LSTM должен использовать последние действительные состояния (игнорировать заполнение) из первого LSTM в качестве исходного состояния.Исходное состояние состоит из cell_state и hidden_state.Но

keras.layers.LSTM(num_units, return_state=True)

возвращает только последние hidden_state и последние cell_state.То, что я хочу, является последним действительным states для каждой последовательности.Я могу использовать Masking, но он не работает с CuDNNLSTM.

Я нашел решение для извлечения cell states для каждого временного шага в Керасе.Но я не смог воспроизвести его в Tensorflow.Утомительный while_loop убивает меня.Я получил ошибку:

ValueError: Initializer for variable while_18/lstm_6/kernel/ is from inside a control-flow construct, such as a loop or conditional. When creating a variable inside a loop or conditional, use a lambda as the initializer.

Ниже вы можете найти мою попытку:

import keras
from keras.layers import Input, LSTM, Lambda
import tensorflow as tf

maxlen = 10
input_dim = 10
units = 5

inputs = Input((maxlen, input_dim), dtype = tf.float32)

rnn = LSTM(units, return_state=True)

def get_indexer(t):
    return Lambda(lambda x, t: x[:, t, :], arguments={'t':t}, output_shape=lambda s: (s[0], s[2]))

def expand(x):
    return keras.backend.expand_dims(x, 1)

expand_layer = Lambda(expand, output_shape=lambda s: (s[0], 1, s[1]))
state = tf.Variable(tf.zeros([10]))
states = tf.TensorArray(dtype=tf.float32, size=0, dynamic_size=True, name='states')
iters = tf.constant(10, name='iters')

def cond(i, iters, states):
    return tf.less(i, iters)

def body(i, iters, states):
    input_t = get_indexer(i)(inputs)  # basically input_t = inputs[:, t, :]
    input_t = expand_layer(input_t)
    output_t, h, c = rnn(input_t, initial_state=state)
    temp_state = h, c
    assign_op = tf.assign(state, temp_state)
    states = states.write(step, state)
    return states

states = tf.while_loop(cond, body, [0, iters, states])