Repository: rwth-i6/returnn

Issue:
CompareLayer: using sources from inside and outside of a rec loop during search
Using the following network:

```
network = {
      "data_int": {"class": "cast", "from": "data", "dtype": "int32"},
      "output": {"class": "rec", "from": "data", "unit": {
        "prev_out": {
          "class": "reinterpret_data", "from": "prev:output", "set_sparse": False},
        "cond": {
          "class": "compare", "from": ["prev_out", "base:data_int"], "kind": "less_equal"},
        "switch": {"class": "switch", "condition": "cond", "true_from": 0.0, "false_from": 0.0},
        "output_prob": {"class": "softmax", "from": "switch", "target": "classes", "loss": "ce"},
        "output": {
          "class": "choice", "from": "output_prob", "beam_size": 3, "target": "classes",
          "input_type": "prob", "initial_output": 0}
      }}
    }
```

the following error is thrown during search:

```
layer <network via test_CompareLayer_in_loop_and_out_loop_sources>/'data' output: Data{'data', [B,T|'time'[B],F|F'feature:data'(5)]}
layer <network via test_CompareLayer_in_loop_and_out_loop_sources>/'data_int' output: Data{'data_int_output', [B,T|'time'[B],F|F'feature:data'(5)], dtype='int32'}
layer <network via test_CompareLayer_in_loop_and_out_loop_sources>/'output' output: Data{'output_output', [T|'time'[B&Beam{'output/output'}(3)],B&Beam{'output/output'}(3)], dtype='int32', sparse_dim=DimensionTag{F'classes:sparse-dim'(5)}}
Rec layer 'output' (search True, train False) sub net:
  Input layers moved out of loop: (#: 0)
    None
  Output layers moved out of loop: (#: 0)
    None
  Layers in loop: (#: 5)
    output
    output_prob
    switch
    cond
    prev_out
  Unused layers: (#: 0)
    None
layer <network via test_CompareLayer_in_loop_and_out_loop_sources>/output(rec-subnet)/'prev_out' output: Data{'prev_out_output', [B&Beam{'output/prev:output'}(3)], dtype='int32', ctx=loop('time'[B])}
layer <network via test_CompareLayer_in_loop_and_out_loop_sources>/output(rec-subnet)/'cond' output: Data{'cond_output', [B&Beam{'output/prev:output'}(3),T|'time'[B&Beam{'output/prev:output'}(3)],F|F'feature:data'(5)], dtype='bool', ctx=loop('time'[B])}

ERROR: Got exception during in-loop construction of layer 'switch':
AssertionError: 

Exception creating layer <network via test_CompareLayer_in_loop_and_out_loop_sources>/output(rec-subnet)/'switch' of class SwitchLayer with opts:
output: <_TemplateLayer(ChoiceLayer)(:template:choice) output/'output' out_type=Data{[B&Beam{'output/output'}(3)], dtype='int32', sparse_dim=DimensionTag{F'classes:sparse-dim'(5)}, ctx=loop('time'[B])} (construction stack None)>
output_prob: <_TemplateLayer(SoftmaxLayer)(:template:softmax) output/'output_prob' out_type=Data{[B,T|'time'[B],F|F'feature:data'(5)], ctx=loop('time'[B])} (construction stack 'output')>
switch: <_TemplateLayer(SwitchLayer)(:template:switch) output/'switch' out_type=Data{[B,T|'time'[B],F|F'feature:data'(5)], ctx=loop('time'[B])} (construction stack 'output_prob')>
{'_name': 'switch',
cond: <_TemplateLayer(CompareLayer)(:template:compare) output/'cond' out_type=Data{[B&Beam{'output/prev:output'}(3),T|'time'[B&Beam{'output/prev:output'}(3)],F|F'feature:data'(5)], dtype='bool', ctx=loop('time'[B])} (construction stack 'switch')>
 '_network': <TFNetwork '<network via test_CompareLayer_in_loop_and_out_loop_sources>/output(rec-subnet)' parent_layer=<RecLayer 'output' out_type=Data{[T|'time'[B&Beam{'output/output'}(3)],B&Beam{'output/output'}(3)], dtype='int32', sparse_dim=DimensionTag{F'classes:sparse-dim'(5)}}> train=False search>,
prev_out: <_TemplateLayer(ReinterpretDataLayer)(:template:reinterpret_data) output/'prev_out' out_type=Data{[B&Beam{'output/prev:output'}(3)], dtype='int32', ctx=loop('time'[B])} (construction stack 'cond')>

 'condition': <CompareLayer output/'cond' out_type=Data{[B&Beam{'output/prev:output'}(3),T|'time'[B&Beam{'output/prev:output'}(3)],F|F'feature:data'(5)], dtype='bool', ctx=loop('time'[B])}>,
Collected (unique) exceptions during template construction:
 'false_from': 0.0,
(Note that many of these can be ignored, or are expected.)
 'name': 'switch',

 'network': <TFNetwork '<network via test_CompareLayer_in_loop_and_out_loop_sources>/output(rec-subnet)' parent_layer=<RecLayer 'output' out_type=Data{[T|'time'[B&Beam{'output/output'}(3)],B&Beam{'output/output'}(3)], dtype='int32', sparse_dim=DimensionTag{F'classes:sparse-dim'(5)}}> train=False search>,
 'sources': [],
EXCEPTION 'true_from': 0.0}

Template network (check out types / shapes):
Template network (check out types / shapes):
Template network (check out types / shapes):
Traceback (most recent call last):
Exception creating layer <network via test_CompareLayer_in_loop_and_out_loop_sources>/'output' of class RecLayer with opts:
  {'_name': 'output',
File '_network': <TFNetwork '<network via test_CompareLayer_in_loop_and_out_loop_sources>' train=False search>,
  'axis': DimensionTag{'time'[B]},
"/mnt/projects/i6/returnn/returnn/tf/layers/ 'n_out': <class 'returnn.util.basic.NotSpecified'>,
rec.py 'name': 'output',
 'network': <TFNetwork '<network via test_CompareLayer_in_loop_and_out_loop_sources>' train=False search>,
" 'output': Data{'output_output', [T|'time'[B&Beam{'output/output'}(3)],B&Beam{'output/output'}(3)], dtype='int32', sparse_dim=DimensionTag{F'classes:sparse-dim'(5)}},
,  'sources': [<SourceLayer 'data' out_type=Data{[B,T|'time'[B],F|F'feature:data'(5)]}>],
line 'unit': <_SubnetworkRecCell '<network via test_CompareLayer_in_loop_and_out_loop_sources>/output(rec-subnet)'>}
 1228, in _SubnetworkRecCell._construct_template.<locals>.GetLayer.construct
    line: self.__call__(layer_name_)
    locals:
      self = <local> <RecLayer construct template GetLayer>(allow_uninitialized_template False, parents None)
      self.__call__ = <local> <bound method _SubnetworkRecCell._construct_template.<locals>.GetLayer.__call__ of <RecLayer construct template GetLayer>(allow_uninitialized_template False, parents None)>
      layer_name_ = <local> 'switch', len = 6
  File "/mnt/projects/i6/returnn/returnn/tf/layers/rec.py", line 1371, in _SubnetworkRecCell._construct_template.<locals>.GetLayer.__call__
    line: res_layer = self.net.construct_layer(
            net_dict=self.net_dict, name=name,
            get_layer=default_get_layer, add_layer=default_get_layer.add_templated_layer)
    locals:
      res_layer = <local> None
      self = <local> <_SubnetworkRecCell '<network via test_CompareLayer_in_loop_and_out_loop_sources>/output(rec-subnet)'>
      self.net = <local> <TFNetwork '<network via test_CompareLayer_in_loop_and_out_loop_sources>/output(rec-subnet)' parent_net=<TFNetwork '<network via test_CompareLayer_in_loop_and_out_loop_sources>' train=False search> train=False search>
      self.net.construct_layer = <local> <bound method TFNetwork.construct_layer of <TFNetwork '<network via test_CompareLayer_in_loop_and_out_loop_sources>/output(rec-subnet)' parent_net=<TFNetwork '<network via test_CompareLayer_in_loop_and_out_loop_sources>' train=False search> train=False search>>
      net_dict = <not found>
      self.net_dict = <local> {'prev_out': {'class': 'reinterpret_data', 'from': 'prev:output', 'set_sparse': False}, 'cond': {'class': 'compare', 'from': ['prev_out', 'base:data_int'], 'kind': 'less_equal'}, 'switch': {'class': 'switch', 'condition': 'cond', 'true_from': 0.0, 'false_from': 0.0}, 'output_prob': {'class': 'sof...
      name = <local> 'switch', len = 6
      get_layer = <not found>
      default_get_layer = <local> <RecLayer construct template GetLayer>(allow_uninitialized_template False, parents 'switch')
      add_layer = <not found>
      default_get_layer.add_templated_layer = <local> <bound method _SubnetworkRecCell._construct_template.<locals>.GetLayer.add_templated_layer of <RecLayer construct template GetLayer>(allow_uninitialized_template False, parents 'switch')>
  File "/mnt/projects/i6/returnn/returnn/tf/network.py", line 935, in TFNetwork.construct_layer
    line: return add_layer(name=name_with_prefix, layer_class=layer_class, **layer_desc)
    locals:
      add_layer = <local> <bound method _SubnetworkRecCell._construct_template.<locals>.GetLayer.add_templated_layer of <RecLayer construct template GetLayer>(allow_uninitialized_template False, parents 'switch')>
      name = <local> 'switch', len = 6
      name_with_prefix = <local> 'switch', len = 6
      layer_class = <local> <class 'returnn.tf.layers.basic.SwitchLayer'>
      layer_desc = <local> {'condition': <_TemplateLayer(CompareLayer)(:template:compare) output/'cond' out_type=Data{[B&Beam{'output/prev:output'}(3),T|'time'[B&Beam{'output/prev:output'}(3)],F|F'feature:data'(5)], dtype='bool', ctx=loop('time'[B])} (construction stack 'switch')>, 'true_from': 0.0, 'false_from': 0.0, '_ne..., len = 6
  File "/mnt/projects/i6/returnn/returnn/tf/layers/rec.py", line 1258, in _SubnetworkRecCell._construct_template.<locals>.GetLayer.add_templated_layer
    line: layer_.kwargs["output"] = layer_class.get_out_data_from_opts(**layer_desc)
    locals:
      layer_ = <local> <_TemplateLayer(SwitchLayer)(:template:switch) output/'switch' out_type=Data{[B,T|'time'[B],F|F'feature:data'(5)], ctx=loop('time'[B])} (construction stack 'output_prob')>
      layer_.kwargs = <local> {'condition': <_TemplateLayer(CompareLayer)(:template:compare) output/'cond' out_type=Data{[B&Beam{'output/prev:output'}(3),T|'time'[B&Beam{'output/prev:output'}(3)],F|F'feature:data'(5)], dtype='bool', ctx=loop('time'[B])} (construction stack 'switch')>, 'true_from': 0.0, 'false_from': 0.0, '_ne..., len = 9
      layer_class = <local> <class 'returnn.tf.layers.basic.SwitchLayer'>
      layer_class.get_out_data_from_opts = <local> <bound method SwitchLayer.get_out_data_from_opts of <class 'returnn.tf.layers.basic.SwitchLayer'>>
      layer_desc = <local> {'condition': <_TemplateLayer(CompareLayer)(:template:compare) output/'cond' out_type=Data{[B&Beam{'output/prev:output'}(3),T|'time'[B&Beam{'output/prev:output'}(3)],F|F'feature:data'(5)], dtype='bool', ctx=loop('time'[B])} (construction stack 'switch')>, 'true_from': 0.0, 'false_from': 0.0, '_ne..., len = 8
  File "/mnt/projects/i6/returnn/returnn/tf/layers/basic.py", line 7061, in SwitchLayer.get_out_data_from_opts
    line: out.sanity_check()
    locals:
      out = <local> Data{'cond_output', [B&Beam{'output/prev:output'}(3),T|'time'[B],F|F'feature:data'(5)], ctx=loop('time'[B])}
      out.sanity_check = <local> <bound method Data.sanity_check of Data{'cond_output', [B&Beam{'output/prev:output'}(3),T|'time'[B],F|F'feature:data'(5)], ctx=loop('time'[B])}>
  File "/mnt/projects/i6/returnn/returnn/tf/util/data.py", line 1808, in Data.sanity_check
    line: assert tag.batch == self.batch or self.batch.is_broadcast()
    locals:
      tag = <local> DimensionTag{'time'[B]}
      tag.batch = <local> BatchInfo{B}
      self = <local> Data{'cond_output', [B&Beam{'output/prev:output'}(3),T|'time'[B],F|F'feature:data'(5)], ctx=loop('time'[B])}
      self.batch = <local> BatchInfo{B, Beam{'output/prev:output'}(3)}
      self.batch.is_broadcast = <local> <bound method BatchInfo.is_broadcast of BatchInfo{B, Beam{'output/prev:output'}(3)}>
AssertionError


ERROR: Got exception during in-loop construction of layer 'output_prob':
AssertionError: 

output: <_TemplateLayer(ChoiceLayer)(:template:choice) output/'output' out_type=Data{[B&Beam{'output/output'}(3)], dtype='int32', sparse_dim=DimensionTag{F'classes:sparse-dim'(5)}, ctx=loop('time'[B])} (construction stack None)>
output_prob: <_TemplateLayer(SoftmaxLayer)(:template:softmax) output/'output_prob' out_type=Data{[B,T|'time'[B],F|F'feature:data'(5)], ctx=loop('time'[B])} (construction stack 'output')>
switch: <_TemplateLayer(SwitchLayer)(:template:switch) output/'switch' out_type=Data{[B,T|'time'[B],F|F'feature:data'(5)], ctx=loop('time'[B])} (construction stack 'output_prob')>
cond: <_TemplateLayer(CompareLayer)(:template:compare) output/'cond' out_type=Data{[B&Beam{'output/prev:output'}(3),T|'time'[B&Beam{'output/prev:output'}(3)],F|F'feature:data'(5)], dtype='bool', ctx=loop('time'[B])} (construction stack 'switch')>
prev_out: <_TemplateLayer(ReinterpretDataLayer)(:template:reinterpret_data) output/'prev_out' out_type=Data{[B&Beam{'output/prev:output'}(3)], dtype='int32', ctx=loop('time'[B])} (construction stack 'cond')>


ERROR: Got exception during in-loop construction of layer 'output':
AssertionError: 

output: <_TemplateLayer(ChoiceLayer)(:template:choice) output/'output' out_type=Data{[B&Beam{'output/output'}(3)], dtype='int32', sparse_dim=DimensionTag{F'classes:sparse-dim'(5)}, ctx=loop('time'[B])} (construction stack None)>
output_prob: <_TemplateLayer(SoftmaxLayer)(:template:softmax) output/'output_prob' out_type=Data{[B,T|'time'[B],F|F'feature:data'(5)], ctx=loop('time'[B])} (construction stack 'output')>
switch: <_TemplateLayer(SwitchLayer)(:template:switch) output/'switch' out_type=Data{[B,T|'time'[B],F|F'feature:data'(5)], ctx=loop('time'[B])} (construction stack 'output_prob')>
cond: <_TemplateLayer(CompareLayer)(:template:compare) output/'cond' out_type=Data{[B&Beam{'output/prev:output'}(3),T|'time'[B&Beam{'output/prev:output'}(3)],F|F'feature:data'(5)], dtype='bool', ctx=loop('time'[B])} (construction stack 'switch')>
prev_out: <_TemplateLayer(ReinterpretDataLayer)(:template:reinterpret_data) output/'prev_out' out_type=Data{[B&Beam{'output/prev:output'}(3)], dtype='int32', ctx=loop('time'[B])} (construction stack 'cond')>

EXCEPTION
...
File "/mnt/projects/i6/returnn/returnn/tf/layers/basic.py", line 7061, in SwitchLayer.get_out_data_from_opts
    line: out.sanity_check()
    locals:
      out = <local> Data{'cond_output', [B&Beam{'output/prev:output'}(3),T|'time'[B],F|F'feature:data'(5)], ctx=loop('time'[B])}
      out.sanity_check = <local> <bound method Data.sanity_check of Data{'cond_output', [B&Beam{'output/prev:output'}(3),T|'time'[B],F|F'feature:data'(5)], ctx=loop('time'[B])}>
  File "/mnt/projects/i6/returnn/returnn/tf/util/data.py", line 1808, in Data.sanity_check
    line: assert tag.batch == self.batch or self.batch.is_broadcast()
    locals:
      tag = <local> DimensionTag{'time'[B]}
      tag.batch = <local> BatchInfo{B}
      self = <local> Data{'cond_output', [B&Beam{'output/prev:output'}(3),T|'time'[B],F|F'feature:data'(5)], ctx=loop('time'[B])}
      self.batch = <local> BatchInfo{B, Beam{'output/prev:output'}(3)}
      self.batch.is_broadcast = <local> <bound method BatchInfo.is_broadcast of BatchInfo{B, Beam{'output/prev:output'}(3)}>
AssertionError
```

The full stack trace is here: https://gist.github.com/robin-p-schmitt/671e74fabb7788313ffd26fc5243cc53

There is no error during training.

I am not sure whether this error is exclusive to the CompareLayer. To me, it seems like there is a problem when combining layer outputs from outside and inside the loop because the time tag has no beam information.


Discussion:
Normally `SearchChoices.translate_to_common_search_beam` should handle this.
> Normally `SearchChoices.translate_to_common_search_beam` should handle this.

Okay, so I just did some debugging and found out that `SearchChoices.translate_to_common_search_beam` is only called once on `data`. I think the problem is that `SwitchLayer.get_out_data_from_opts` does a sanity check on the output and this happens **before** `SearchChoices.translate_to_common_search_beam` is called. If I remove the sanity check in `SwitchLayer`, the search runs without a problem and `debug_print_layer_output_template` shows, that there now is the previously missing beam information in the time tag. Is the sanity check in `SwitchLayer` needed or could this be left out or replaced by something else?
Yes, you can remove the `sanity_check` there because we anyway do it later again, also after calling `fixup_out_data` (see logic in `TFNetwork._create_layer`).

Please make a PR with the test case and the fix.


Resolving patch:
diff --git a/returnn/tf/layers/basic.py b/returnn/tf/layers/basic.py
--- a/returnn/tf/layers/basic.py
+++ b/returnn/tf/layers/basic.py
@@ -7179,7 +7179,6 @@ def get_source_template(source, source_name):
     out = Data.get_common_data([true_data, false_data, condition.output.copy_template()])
     out.dtype = true_data.dtype
     out.sparse_dim = true_data.sparse_dim
-    out.sanity_check()
     return out.copy(name="%s_output" % name)
 
   def get_dep_layers(self):


Test patch:
diff --git a/tests/test_TFNetworkLayer.py b/tests/test_TFNetworkLayer.py
--- a/tests/test_TFNetworkLayer.py
+++ b/tests/test_TFNetworkLayer.py
@@ -1266,6 +1266,38 @@ def test_CompareLayer_allow_broadcast_all_sources():
     })
 
 
+def test_SwitchLayer_sanity_check():
+  """
+  https://github.com/rwth-i6/returnn/issues/800
+  """
+  from returnn.tf.util.data import DimensionTag
+  with make_scope():
+    n_out = 5
+    time_tag = DimensionTag(DimensionTag.Types.Spatial, "time")
+    config = Config({
+      "debug_print_layer_output_template": True,
+      "extern_data": {
+        "data": {"dim": n_out, "same_dim_tags_as": {"t": time_tag}},
+        "classes": {"dim": n_out, "sparse": True, "same_dim_tags_as": {"t": time_tag}}
+      }})
+    net = TFNetwork(config=config, search_flag=True)
+    net.construct_from_dict({
+      "data_int": {"class": "cast", "from": "data", "dtype": "int32"},
+      "output": {"class": "rec", "from": "data", "unit": {
+        "prev_out": {
+          "class": "reinterpret_data", "from": "prev:output", "set_sparse": False},
+        "cond": {
+          "class": "compare", "from": ["prev_out", "base:data_int"], "kind": "less_equal"},
+        "switch": {"class": "switch", "condition": "cond", "true_from": 0.0, "false_from": 0.0},
+        "switch_red": {"class": "reduce", "from": "switch", "axis": "t", "mode": "max"},
+        "output_prob": {"class": "softmax", "from": "switch_red", "target": "classes", "loss": "ce"},
+        "output": {
+          "class": "choice", "from": "output_prob", "beam_size": 3, "target": "classes",
+          "input_type": "prob", "initial_output": 0}
+      }}
+    })
+
+
 def test_dot_layer_shuffled_remaining_dims_static():
   with make_scope() as session:
     import numpy as np