Repository: rwth-i6/returnn

Issue:
SplitLayer is never optimized as output layer
Consider some net dict of a rec unit like this (see #559 for a full test case):
```
  input: {...., from: "prev:input"},  # cannot be optimized out
  split: {class: split, from: input, size_splits: [3, 7]},
  output: {class: linear, from: ["split/0", "split/1"]}  # can be optimized out
```
Then, `split` should also be optimized out as an output layer.
This is not the case however, because:
`split/0` and `split/1` have `collocate_with=["split"]` set by `_TemplateLayer.get_sub_layer`, but also an dependency on their parent layer `split`.

E.g. `split/0` is not optimized out because the logic of `output_can_move_out(split/0)` does the following:
1. it checks whether the collocated parent `split` can be moved out with `_collocated={split/0}`
2. however, `split` cannot be moved out because `split/1` depends on it. And `split/1` has not been moved out yet. [edited.]

Symmetrically, `split/1` is also not optimized out.

--------------------

(Actually, in my use case, I don't even have a layer `input` that recurrently depends on itself. But instead, I use a `:i` layer, which currently never is optimized out as an input layer:
```python3
      if layer.name in [":i", "end"]:  # currently not fully implemented
        return False
```
Somewhat related - what is missing of the implementation there?)


Discussion:

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
@@ -2901,13 +2901,13 @@ def output_can_move_out(layer, _collocated=None):
       for other_layer in layers_in_loop:
         if other_layer in _collocated:
           continue
-        if layer in other_layer.dependencies:
-          return False
-        if other_layer.name in layer.collocate_with:
+        if other_layer.name in layer.collocate_with or layer.name in other_layer.collocate_with:
           if output_can_move_out(other_layer, _collocated=_collocated):
-            pass  # this is ok, we can move all out
+            continue  # this is ok, we can move all out
           else:
             return False
+        if layer in other_layer.dependencies:
+          return False
       return True
 
     def find_output_layer_to_move_out():


Test patch:
diff --git a/tests/test_TFNetworkRecLayer.py b/tests/test_TFNetworkRecLayer.py
--- a/tests/test_TFNetworkRecLayer.py
+++ b/tests/test_TFNetworkRecLayer.py
@@ -3598,6 +3598,45 @@ def test_reclayer_optimize_out_access_split():
     other_subnet_layers={"split": {"class": "split", "from": ["data:source"], "size_splits": [5, 8]}})
 
 
+def test_SplitLayer_move_out_as_output_layer():
+  with make_scope() as session:
+    config = Config()
+    config.update({
+      "debug_print_layer_output_template": True,
+      "extern_data": {
+        "data": {"dim": 11, "available_for_inference": True},
+        "classes": {"dim": 10, "available_for_inference": False},
+      },
+      "network": {
+        "output": {
+          "class": "rec", "from": [], "max_seq_len": 5, "target": "classes",
+          "unit": {
+            "input": {"class": "linear", "n_out": 10, "from": ["prev:input"], "with_bias": True},  # not optimized out
+            "split": {"class": "split", "from": "input", "size_splits": [1, 9]},  # goal: optimize this out
+            "output": {
+              "class": "linear", "from": ["split/0", "split/1"], "target": "classes"},  # optimized out
+            "end_compare": {"class": "compare", "kind": "greater", "from": ["split/0"], "value": 0},
+            "end": {"class": "squeeze", "from": ["end_compare"], "axis": "F"},
+          }}
+      }
+    })
+    network = TFNetwork(config=config, train_flag=True)
+    network.construct_from_dict(config.typed_dict["network"])
+    rec_layer = network.get_layer("output")
+    assert isinstance(rec_layer, RecLayer)
+    rec_cell = rec_layer.cell
+    assert set(rec_cell.input_layers_moved_out) == set()
+    assert set(rec_cell.output_layers_moved_out) == {"output", "split", "split/0", "split/1"}
+    assert set(rec_cell.layers_in_loop) == {"input"}
+
+    # run but don't care about the result
+    session.run(tf_compat.v1.global_variables_initializer())
+    output_layer = network.get_default_output_layer(must_exist=True)
+    from test_TFNetworkLayer import make_feed_dict
+    feed_dict = make_feed_dict(list(network.extern_data.data.values()))
+    session.run(output_layer.output.placeholder, feed_dict=feed_dict)
+
+
 def test_reclayer_att_with_kv_in_rec():
   net_dict = {
     'decision': {'class': 'decide', 'from': ['output'], 'loss': 'edit_distance', 'loss_opts': {}, 'target': 'classes'},