Repository: rwth-i6/returnn

Issue:
NativeLstm (version 1) does not correctly handle first/last state
This test case fails:
```
def test_reclayer_optimize_out_rec_nativelstm():
  check_reclayer_optimize_out({"class": "rec", "unit": "NativeLstm"})
```
Note that this is specifically about the first version. It works correctly with `NativeLstm2`.
Unfortunately our `unit="lstm"` default currently uses `NativeLstm` internally. So when you just used `lstm`, you have the same problem.
The same problem also for `unit="lstmp"`.
(Maybe it's anyway a good idea to change that default to `NativeLstm2`.)

From briefly looking at it, I think the reason is that `NativeLstm` does not handle the first (initial) and last state correctly or completely. It looks like it only handles the cell state `c` but not `h`.

When the LSTM is not optimized out of the loop, it means this problem occurs every single frame.

This means, when you used sth like this in your config, the LSTM was not behaving correctly:
```
...
output: {class: rec, ..., unit: {
  ...
  s: {class: rec, unit: lstm, ...},
  ...
}}
```
This is quite standard in some attention-based LSTM-based encoder-decoder models.

Note that when you used `rnn_cell`, you could not make use of `NativeLstm` anyway, so you did not have this problem (only the problem that it is slow).
Also, again as stated before, when you use `NativeLstm2`, you also did not have this problem.



Discussion:
@mmz33 @JackTemaki @patrick-wilken @michelwi maybe some of you are affected by this.

I'm not exactly sure what to do about it. Maybe just disallow `NativeLstm` inside another rec layer so that people do not run into this. Maybe additionally changing the `RecLayer` `unit="lstm"` default to `NativeLstm2`.
I think I started to always use nativelstm2 some time ago, there should be no remaining `nativelstm` layers in my setups, but I will verify. Some seutps transfered from other people might have it.
Thanks for the heads up, I was not using `NativeLstm` (v1) in my experiments.

Resolving patch:
diff --git a/returnn/tf/layers/rec.py b/returnn/tf/layers/rec.py
--- a/returnn/tf/layers/rec.py
+++ b/returnn/tf/layers/rec.py
@@ -83,7 +83,7 @@ class RecLayer(_ConcatInputLayer):
 
   layer_class = "rec"
   recurrent = True
-  _default_lstm_unit = "nativelstm"  # TFNativeOp.NativeLstmCell
+  _default_lstm_unit = "NativeLSTM2"  # TFNativeOp.NativeLstmCell
   SubnetworkRecCell = None  # type: typing.Type[_SubnetworkRecCell]  # set below
 
   def __init__(self,
@@ -971,12 +971,21 @@ def _get_output_cudnn(self, cell):
 
   def _get_output_native_rec_op(self, cell):
     """
-    :param TFNativeOp.RecSeqCellOp cell:
+    :param returnn.tf.native_op.RecSeqCellOp cell:
     :return: output of shape (time, batch, dim)
     :rtype: tf.Tensor
     """
     from returnn.tf.util.basic import dot, sequence_mask_time_major, directed, to_int32_64, set_param_axes_split_info
 
+    # Some error checking.
+    from returnn.tf import native_op
+    if isinstance(cell, native_op.NativeLstmCell):
+      if self._initial_state is not None:
+        raise Exception(
+          "%s: NativeLstm (v1) has broken (only partial) support for initial_state. "
+          "See here for details: https://github.com/rwth-i6/returnn/issues/813. "
+          "Use 'NativeLstm2' instead." % self)
+
     assert self._max_seq_len is None
     assert self.input_data
     in_data, x, seq_len = self._get_input()


Test patch:
diff --git a/tests/test_TFNetworkRecLayer.py b/tests/test_TFNetworkRecLayer.py
--- a/tests/test_TFNetworkRecLayer.py
+++ b/tests/test_TFNetworkRecLayer.py
@@ -4029,6 +4029,42 @@ def test_SplitLayer_move_out_as_output_layer():
     session.run(output_layer.output.placeholder, feed_dict=feed_dict)
 
 
+def test_reclayer_inner_nativelstm1():
+  # https://github.com/rwth-i6/returnn/issues/813
+  # NativeLstm1 is broken inside RecLayer.
+  # This should cause an exception. We test this here.
+  def _test_inner_rec_unit(unit):
+    print("*** test rec unit", unit)
+    net_dict = {
+      "output": {
+        "class": "rec", "from": "data",
+        "optimize_move_layers_out": False,  # make sure it stays inside the loop, for the actual test
+        "unit": {
+          "output": {"class": "rec", "unit": unit, "from": "data:source", "n_out": 7}
+        }
+      }
+    }
+    with make_scope() as session:
+      config = Config({"extern_data": {"data": {"dim": 3}}})
+      net = TFNetwork(config=config)
+      net.construct_from_dict(net_dict)
+      net.initialize_params(session)
+      out = net.get_default_output_layer()
+      from test_TFNetworkLayer import make_feed_dict
+      session.run(out.output.placeholder, feed_dict=make_feed_dict(net.extern_data))
+
+  for unit in ["NativeLstm2", "lstm"]:
+    _test_inner_rec_unit(unit)  # should be no error
+  for unit in ["NativeLstm"]:
+    try:
+      _test_inner_rec_unit(unit)
+    except Exception as exc:
+      print("Expected exception:", exc)
+      assert "https://github.com/rwth-i6/returnn/issues/813" in str(exc)
+    else:
+      raise Exception("Expect to get exception for unit %r. https://github.com/rwth-i6/returnn/issues/813" % unit)
+
+
 def test_reclayer_att_with_kv_in_rec():
   net_dict = {
     'decision': {'class': 'decide', 'from': ['output'], 'loss': 'edit_distance', 'loss_opts': {}, 'target': 'classes'},