Repository: rwth-i6/returnn

Issue:
MaskedComputationLayer: dyn_size error
The following test case leads to an error:

```
def test_MaskedComputationLayer_dyn_size_none():
    with make_scope() as session:
      config = Config({"debug_print_layer_output_template": True})
      net = TFNetwork(
        extern_data=ExternData(
          {"data": {"dim": 20, "sparse": False}, "classes": {"dim": 20, "sparse": True}}),
        config=config, train_flag=False, search_flag=True)
      net_dict = {
        "rec_loop": {
          "class": "rec", "from": "data", "unit": {
            "lin": {"class": "linear", "activation": "softmax", "from": "prev:output", "n_out": 20},
            "output": {
              "class": "choice", "from": "lin", "beam_size": 4, "target": "classes", "initial_output": 0
            }
          }},
        "data_red": {"class": "reduce", "mode": "mean", "from": "data", "axis": "f"},
        "mask": {"class": "compare", "from": "data_red", "value": 0.0, "kind": "greater"},
        "masked_data": {
          "class": "masked_computation", "from": "rec_loop", "mask": "mask", "unit": {
            "class": "copy", "from": "data"
          }
        },
        "output": {"class": "decide", "from": "masked_data"}
      }
      net.construct_from_dict(net_dict)
```

The error is the following:

```
Exception creating layer <network via test_MaskedComputationLayer_dyn_size_none>/'masked_data' of class MaskedComputationLayer with opts:
{'_layer_class': <class 'returnn.tf.layers.basic.CopyLayer'>,
 '_layer_desc': {'_name': 'masked_data',
                 '_network': <TFNetwork '<network via test_MaskedComputationLayer_dyn_size_none>(extra._internal_template.masked(masked_data))' train=False search>,
                 'encapsulate': True,
                 'name': 'masked_data',
                 'network': <TFNetwork '<network via test_MaskedComputationLayer_dyn_size_none>(extra._internal_template.masked(masked_data))' train=False search>,
                 'sources': [<WrappedInternalLayer 'rec_loop' out_type=Data{[T|'masked_data:masked:time'[?],B&Beam{'rec_loop/output'}(4)], dtype='int32', sparse_dim=Dim{F'classes:sparse-dim'(20)}}>]},
 '_name': 'masked_data',
 '_network': <TFNetwork '<network via test_MaskedComputationLayer_dyn_size_none>' train=False search>,
 '_parent_layer_cache': {},
 '_src_common_search_choices': <SearchChoices owner='output' beam_size=4 beam_scores=shaped:(None,None)>,
 'in_spatial_dim': Dim{'time:var:extern_data:data'[B&Beam{'rec_loop/output'}(4)]},
 'mask': <SelectSearchSourcesLayer 'mask' <SearchChoices owner='output' beam_size=4 beam_scores=shaped:(None,None)> out_type=Data{[B&Beam{'rec_loop/output'}(4),T|'time:var:extern_data:data'[B&Beam{'rec_loop/output'}(4)]], dtype='bool'}>,
 'masked_from': None,
 'name': 'masked_data',
 'network': <TFNetwork '<network via test_MaskedComputationLayer_dyn_size_none>' train=False search>,
 'out_spatial_dim': Dim{'masked_data:masked:time'[B]},
 'output': Data{'masked_data_output', [T|'masked_data:masked:time'[?],B&Beam{'rec_loop/output'}(4)], dtype='int32', sparse_dim=Dim{F'classes:sparse-dim'(20)}},
 'sources': [<RecLayer 'rec_loop' out_type=Data{[T|'time:var:extern_data:data'[B&Beam{'rec_loop/output'}(4)],B&Beam{'rec_loop/output'}(4)], dtype='int32', sparse_dim=Dim{F'classes:sparse-dim'(20)}}>],
 'unit': {'class': 'copy', 'from': 'data'}}
EXCEPTION
...
  File "/mnt/projects/i6/returnn/returnn/tf/layers/rec.py", line 7752, in MaskedComputationLayer.__init__
    line: self.sub_layer.output.sanity_check()
    locals:
      self = <local> <MaskedComputationLayer 'masked_data' out_type=Data{[T|'masked_data:masked:time'[?],B&Beam{'rec_loop/output'}(4)], dtype='int32', sparse_dim=Dim{F'classes:sparse-dim'(20)}}>
      self.sub_layer = <local> <CopyLayer 'masked_data' out_type=Data{[T|'masked_data:masked:time'[?],B&Beam{'rec_loop/output'}(4)], dtype='int32', sparse_dim=Dim{F'classes:sparse-dim'(20)}}>
      self.sub_layer.output = <local> Data{'masked_data_output', [T|'masked_data:masked:time'[?],B&Beam{'rec_loop/output'}(4)], dtype='int32', sparse_dim=Dim{F'classes:sparse-dim'(20)}}
      self.sub_layer.output.sanity_check = <local> <bound method Data.sanity_check of Data{'masked_data_output', [T|'masked_data:masked:time'[?],B&Beam{'rec_loop/output'}(4)], dtype='int32', sparse_dim=Dim{F'classes:sparse-dim'(20)}}>
  File "/mnt/projects/i6/returnn/returnn/tf/util/data.py", line 2879, in Data.sanity_check
    line: assert tag.dyn_size is not None
    locals:
      tag = <local> Dim{'masked_data:masked:time'[?]}
      tag.dyn_size = <local> None
AssertionError
```

The full log is here: https://gist.github.com/robin-p-schmitt/981bd04d967f212aee8c7444aab373d3


Discussion:
This looks related to search. The error also does not occur without search in the test.
I think the error is because `mask` is without a beam but `rec_loop` is with a beam.
Is this actually the case you have? Do you have `mask` without beam and `from` with beam?
Is this correct? Why is your mask without beam?
**Edit** Ok, the error also happens when `mask` is also with a beam. I extended the test. This is probably the more common case.

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
@@ -7653,7 +7653,7 @@ def __init__(self, mask, unit, masked_from,
           new_size = masked_from.output.get_sequence_lengths()
         else:
           new_size = idxs[-1]  # [B]
-          out_spatial_dim.dyn_size = new_size
+          out_spatial_dim.get_for_batch_ctx(self.output.batch, self.output.control_flow_ctx).dyn_size = new_size
         new_time = tf.reduce_max(new_size)  # T'
         idxs = where_bc(mask_t, idxs - 1, new_time)
 
@@ -7744,6 +7744,7 @@ def sub_get_layer(sub_layer_name):
       assert out_spatial_dim
       axis = sub_out_template.get_axis_from_description(in_spatial_dim)
       sub_out_template = sub_out_template.copy_template_replace_dim_tag(axis=axis, new_dim_tag=out_spatial_dim)
+    sub_out_template = layer_class.fixup_out_data(output=sub_out_template, **layer_desc)
     layer_desc["output"] = sub_out_template
 
     self.sub_layer = layer_class(**layer_desc)


Test patch:
diff --git a/tests/test_TFNetworkRecLayer.py b/tests/test_TFNetworkRecLayer.py
--- a/tests/test_TFNetworkRecLayer.py
+++ b/tests/test_TFNetworkRecLayer.py
@@ -7034,6 +7034,35 @@ def test_MaskedComputationLayer_rec_subnet_name_scope():
     assert_equal(set(p.name for p in params), {"linear/W:0", "linear/b:0"})
 
 
+def test_MaskedComputationLayer_dyn_size_none():
+  # https://github.com/rwth-i6/returnn/issues/1008
+  with make_scope() as session:
+    config = Config()
+    net = TFNetwork(
+      extern_data=ExternData(
+        {"data": {"dim": 20, "sparse": False}, "classes": {"dim": 20, "sparse": True}}),
+      config=config,
+      train_flag=False,
+      search_flag=True)
+    net_dict = {
+      "rec_loop": {
+        "class": "rec", "from": "data", "unit": {
+          "lin": {"class": "linear", "activation": "softmax", "from": "prev:output", "n_out": 20},
+          "output": {
+            "class": "choice", "from": "lin", "beam_size": 4, "target": "classes", "initial_output": 0
+          }
+        }},
+      "data_red": {"class": "reduce", "mode": "mean", "from": "data", "axis": "f"},
+      "mask": {"class": "compare", "from": "rec_loop", "value": 10, "kind": "greater"},
+      "masked_data": {
+        "class": "masked_computation", "from": "rec_loop", "mask": "mask",
+        "unit": {"class": "copy", "from": "data"}
+      },
+      "output": {"class": "decide", "from": "masked_data"}
+    }
+    net.construct_from_dict(net_dict)
+
+
 def test_subnet_deps_search():
   beam_size = 3
   EncKeyTotalDim = 10