Как я могу добавить состояние ячейки в LSTM / RNN Tensorflow?

Question

Я хотел бы добавить дополнительное (двоичное) состояние ячейки к моей модели Tensorflow LSTM.

Поэтому я пытаюсь настроить следующую шаговую функцию LSTM:

 def step(cell_inputs, cell_states):

  h_tm1 = cell_states[0]  # previous memory state
  c_tm1 = cell_states[1]  # previous carry state

  z = K.dot(cell_inputs, kernel)
  z += K.dot(h_tm1, recurrent_kernel)
  z = K.bias_add(z, bias)

  z0, z1, z2, z3 = array_ops.split(z, 4, axis=1)

  i = recurrent_activation(z0)
  f = recurrent_activation(z1)
  c = f * c_tm1 + i * activation(z2)
  o = recurrent_activation(z3)

  h = o * activation(c)
  return h, [h, c]

Эта шаговая функция будет использоваться в серверной части RNN Keras:

last_output, outputs, new_states = K.rnn(
    step,
    inputs, [init_h, init_c],
    constants=None,
    unroll=False,
    time_major=time_major,
    mask=mask,
    go_backwards=go_backwards,
    input_length=sequence_lengths
    if sequence_lengths is not None else timesteps)

Мое дополнительное состояние ячейки b устанавливается на 0 или 1 на каждом слое и временном шаге. Я должен иметь доступ к b(layer-1,timestep), а также b(layer,timestep-1). Как это можно реализовать? Нужно ли добавлять состояние ячейки b в переменную cell_inputs, чтобы я мог прочитать b(layer-1,timestep)? И нужно ли мне добавить b к переменной cell_states, чтобы я мог прочитать b(layer,timestep-1)?

Любая помощь очень ценится!