Repository: rwth-i6/returnn

Issue:
Passing recurrent state (per frame) to layer, not just initial state
`initial_state` defines the initial state (state in first iteration of the loop). It can point to other layers.
When this is another layer inside the loop, this fails (which is correct).

Example:
```
{'loop': {'class': 'rec',
          'from': [],
          'name_scope': '',
          'unit': {'lstm': {'class': 'rec',
                            'from': 'rec_unstack',
                            'initial_state': {'c': 'prev:state.lstm.c',
                                              'h': 'prev:state.lstm.h'},
                            'n_out': 13,
                            'unit': 'nativelstm2'},
                   'lstm_state_c': {'class': 'get_last_hidden_state',
                                    'from': 'lstm',
                                    'key': 'c',
                                    'n_out': 13},
                   'lstm_state_h': {'class': 'get_last_hidden_state',
                                    'from': 'lstm',
                                    'key': 'h',
                                    'n_out': 13},
                   'output': {'class': 'copy', 'from': 'lstm'},
                   'rec_unstack': {'axis': 'T',
                                   'class': 'rec_unstack',
                                   'from': 'base:data:data'},
                   'state.lstm.c': {'class': 'copy', 'from': 'lstm_state_c'},
                   'state.lstm.h': {'class': 'copy', 'from': 'lstm_state_h'}}},
 'output': {'class': 'copy', 'from': 'loop/output'}}
```
This example makes the hidden state logic of the inner `lstm` layer kind of redundant but explicit, and you can imagine that other variations are possible.
The problem is `initial_state` pointing to eg `prev:state.lstm.c`, although the `prev:` part is not even relevant. This is the (prev) template layer, which is not initialized when it is first used via `get_init_loop_vars`:
```
...
  File "/Users/az/Programmierung/returnn/returnn/tf/network.py", line 1000, in _create_layer
    layer = layer_class(**layer_desc)
  File "/Users/az/Programmierung/returnn/returnn/tf/layers/rec.py", line 239, in __init__
    y = self._get_output_subnet_unit(self.cell)
  File "/Users/az/Programmierung/returnn/returnn/tf/layers/rec.py", line 930, in _get_output_subnet_unit
    output = cell.get_output()
  File "/Users/az/Programmierung/returnn/returnn/tf/layers/rec.py", line 2675, in get_output
    self.get_init_loop_vars(),
  File "/Users/az/Programmierung/returnn/returnn/tf/layers/rec.py", line 1824, in get_init_loop_vars
    self._initial_extra_outputs = {
  File "/Users/az/Programmierung/returnn/returnn/tf/layers/rec.py", line 1825, in <dictcomp>
    k: self._get_init_extra_outputs(k)
  File "/Users/az/Programmierung/returnn/returnn/tf/layers/rec.py", line 1799, in _get_init_extra_outputs
    d = cl.get_rec_initial_extra_outputs(
  File "/Users/az/Programmierung/returnn/returnn/tf/layers/rec.py", line 449, in get_rec_initial_extra_outputs
    return {"state": RnnCellLayer.get_rec_initial_state(**kwargs)}
  File "/Users/az/Programmierung/returnn/returnn/tf/layers/rec.py", line 4335, in get_rec_initial_state
    s = {k: cls.get_rec_initial_state_inner(initial_shape=(batch_dim, d), initial_state=init_value[k], key=k,
  File "/Users/az/Programmierung/returnn/returnn/tf/layers/rec.py", line 4335, in <dictcomp>
    s = {k: cls.get_rec_initial_state_inner(initial_shape=(batch_dim, d), initial_state=init_value[k], key=k,
  File "/Users/az/Programmierung/returnn/returnn/tf/layers/rec.py", line 4380, in get_rec_initial_state_inner
    assert initial_state.output.placeholder is not None
AssertionError
```
(Caught explicitly via #731.)

This is now also more a question of the expected behavior, or how we want it to be.

The intention was that we explicitly pass the state to the `lstm` layer and actually not use the hidden state. For the case that the `lstm` layer is inside the loop. Which is anyway enforced here due to the circular dependency. Otherwise we would expect that it should be correct with automatic optimization.

The main problem is that those state vars ("extra output") which depends on other layers inside the loop will make the state var obsolete.

But also, the semantics of `initial_state` do not really match the intention here. If you set it to e.g. 0, then this will really only be the initial state and subsequent states will not be 0. So it will not pass this state every loop frame but really only in the initial frame.

So it means we need a different mechanism here. Maybe some argument `state`, which is very similar to `initial_state` but which defines the input of state per frame.


Discussion:

Resolving patch:
diff --git a/returnn/tf/layers/base.py b/returnn/tf/layers/base.py
--- a/returnn/tf/layers/base.py
+++ b/returnn/tf/layers/base.py
@@ -67,6 +67,7 @@ def __init__(self, name, network, output, n_out=NotSpecified, out_type=None, sou
                param_variational_noise=None,
                updater_opts=None,
                initial_output=None,
+               state=None,
                rec_previous_layer=None,
                collocate_with=None,
                trainable=None,
@@ -117,6 +118,7 @@ def __init__(self, name, network, output, n_out=NotSpecified, out_type=None, sou
     :param int|None copy_output_loss_from_source_idx: if set, will copy output_loss from this source
     :param bool|dict batch_norm: see self.batch_norm()
     :param str|float initial_output: used for recurrent layer, see self.get_rec_initial_output()
+    :param state: explicitly defines the rec state. initial_state would define the initial state (in the first frame)
     :param LayerBase|None rec_previous_layer: via the recurrent layer, layer (template) which represents the past of us.
       You would not explicitly set this in a config. This is automatically, internally, via :class:`RecLayer`.
     :param list[str]|None collocate_with: in the rec layer, collocate with the specified other layers
@@ -190,6 +192,30 @@ def __init__(self, name, network, output, n_out=NotSpecified, out_type=None, sou
     self.register_as_extern_data = register_as_extern_data
     # Stats will be collected by the engine.
     self.stats = {}  # type: typing.Dict[str,tf.Tensor]
+    self._set_prev_state(state)
+
+  def _set_prev_state(self, state):
+    if state is None:
+      return
+    assert self._rec_previous_layer, "%s: cannot pass 'state' %r when not inside a rec loop" % (self, state)
+    from tensorflow.python.util import nest
+
+    def _map_to_state_tensor(orig_state, state_layer):
+      assert isinstance(orig_state, tf.Tensor)
+      if state_layer is None:  # this is allowed, can be partial
+        return orig_state
+      assert isinstance(state_layer, LayerBase)
+      assert orig_state.shape.as_list() == list(state_layer.output.batch_shape)
+      return state_layer.output.placeholder
+
+    if set(self._rec_previous_layer.rec_vars_outputs.keys()) == {"state"}:
+      rec_prev_layer_state = self._rec_previous_layer.rec_vars_outputs["state"]
+      nest.assert_same_structure(rec_prev_layer_state, state)
+      self._rec_previous_layer.rec_vars_outputs["state"] = nest.map_structure(
+        _map_to_state_tensor, rec_prev_layer_state, state)
+      return
+    raise NotImplementedError(
+      "%s: explicit 'state' %r, internal states %r" % (self, state, self._rec_previous_layer.rec_vars_outputs))
 
   def post_init(self, layer_desc):
     """
@@ -550,6 +576,9 @@ def transform_config_dict(cls, d, network, get_layer):
     rec_previous_layer = root_ctx_net.layers.get("prev:%s%s" % (prefix, d["_name"]))
     if rec_previous_layer:
       d["rec_previous_layer"] = rec_previous_layer
+    if d.get("state", None) is not None:
+      from tensorflow.python.util import nest
+      d["state"] = nest.map_structure(get_layer, d["state"])
 
   @classmethod
   def _guess_n_out_from_target_and_opt_loss(cls, network, target, target_layers, loss_class_name, get_layer):


Test patch: