Repository: rwth-i6/returnn

Issue:
ReinterpretDataLayer: beam information of new DimensionTag is not set correctly
When using the following network for search:

```
network = {
      "data_new": {"class": "reinterpret_data", "from": "data",
                   "set_dim_tags": {"t": new_time_tag}},
      "output": {"class": "rec", "from": "data", "unit": {
        "prev_out0": {
          "class": "reinterpret_data", "from": "prev:output", "set_sparse": False},
        "prev_out1": {"class": "cast", "from": "prev_out0", "dtype": "float32"},
        "prev_out": {"class": "expand_dims", "from": "prev_out1", "axis": "f"},
        "data_concat": {
          "class": "copy", "from": ["base:data_new", "prev_out"]
        },
        "data_red": {"class": "reduce", "from": "data_concat", "axis": "stag:new-time", "mode": "max"},
        "output_prob": {"class": "softmax", "from": "data_red", "target": "classes", "loss": "ce"},
        "output": {
          "class": "choice", "from": "output_prob", "beam_size": 3, "target": "classes",
          "input_type": "prob", "initial_output": 0}
      }}
    }
```

with the following data:

```
    n_out = 5
    time_tag = DimensionTag(DimensionTag.Types.Spatial, "time")
    new_time_tag = DimensionTag(DimensionTag.Types.Spatial, "new-time")
    config = Config({
      "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}}
      }})
```

the following error occurs:


```
layer <network via test_concat_new_dim_tag>/'data' output: Data{'data', [B,T|'time'[B],F|F'feature:data'(5)]}
layer <network via test_concat_new_dim_tag>/'data_new' output: Data{'data_new_output', [B,T|'new-time'[?],F|F'feature:data'(5)]}
layer <network via test_concat_new_dim_tag>/'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: (#: 7)
    output
    output_prob
    data_red
    data_concat
    prev_out
    prev_out1
    prev_out0
  Unused layers: (#: 0)
    None
layer <network via test_concat_new_dim_tag>/output(rec-subnet)/'prev_out0' output: Data{'prev_out0_output', [B&Beam{'output/prev:output'}(3)], dtype='int32', ctx=loop('time'[B])}
layer <network via test_concat_new_dim_tag>/output(rec-subnet)/'prev_out1' output: Data{'prev_out1_output', [B&Beam{'output/prev:output'}(3)], ctx=loop('time'[B])}
layer <network via test_concat_new_dim_tag>/output(rec-subnet)/'prev_out' output: Data{'prev_out_output', [B&Beam{'output/prev:output'}(3),F|'prev_out_expand_dims'(1)], ctx=loop('time'[B])}
layer <network via test_concat_new_dim_tag>/output(rec-subnet)/'data_concat' output: Data{'data_concat_output', [B&Beam{'output/prev:output'}(3),T|'new-time'[B&Beam{'output/prev:output'}(3)],F|F'data_concat_output_feature'(6)], ctx=loop('time'[B])}
Exception creating layer <network via test_concat_new_dim_tag>/output(rec-subnet)/'data_concat' of class CopyLayer with opts:
{'_name': 'data_concat',
 '_network': <TFNetwork '<network via test_concat_new_dim_tag>/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>,
 '_src_common_search_choices': <SearchChoices owner='prev:output' beam_size=3 beam_scores=shaped:(None,None)>,
 'name': 'data_concat',
 'network': <TFNetwork '<network via test_concat_new_dim_tag>/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>,
 'output': Data{'data_concat_output', [B&Beam{'output/prev:output'}(3),T|'new-time'[B&Beam{'output/prev:output'}(3)],F|F'data_concat_output_feature'(6)], ctx=loop('time'[B])},
 'sources': [<SelectSearchSourcesLayer 'data_new' <SearchChoices owner='prev:output' beam_size=3 beam_scores=shaped:(None,None)> out_type=Data{[B&Beam{'output/prev:output'}(3),T|'new-time'[B&Beam{'output/prev:output'}(3)],F|F'feature:data'(5)]}>,
             <ExpandDimsLayer output/'prev_out' out_type=Data{[B&Beam{'output/prev:output'}(3),F|'prev_out_expand_dims'(1)], ctx=loop('time'[B])}>]}
Exception creating layer <network via test_concat_new_dim_tag>/'output' of class RecLayer with opts:
{'_name': 'output',
 '_network': <TFNetwork '<network via test_concat_new_dim_tag>' train=False search>,
 'axis': DimensionTag{'time'[B]},
 'n_out': <class 'returnn.util.basic.NotSpecified'>,
 'name': 'output',
 'network': <TFNetwork '<network via test_concat_new_dim_tag>' 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)]}>],
 'unit': <_SubnetworkRecCell '<network via test_concat_new_dim_tag>/output(rec-subnet)'>}

ERROR: Got exception during in-loop construction of layer 'data_concat':
AssertionError: 

Template network (check out types / shapes):
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&Beam{'output/prev:output'}(3),F|F'output_prob:feature-dense'(5)], ctx=loop('time'[B])} (construction stack 'output')>
data_red: <_TemplateLayer(ReduceLayer)(:template:reduce) output/'data_red' out_type=Data{[B&Beam{'output/prev:output'}(3),F|F'data_concat_output_feature'(6)], ctx=loop('time'[B])} (construction stack 'output_prob')>
data_concat: <_TemplateLayer(CopyLayer)(:template:copy) output/'data_concat' out_type=Data{[B&Beam{'output/prev:output'}(3),T|'new-time'[B&Beam{'output/prev:output'}(3)],F|F'data_concat_output_feature'(6)], ctx=loop('time'[B])} (construction stack 'data_red')>
prev_out: <_TemplateLayer(ExpandDimsLayer)(:template:expand_dims) output/'prev_out' out_type=Data{[B&Beam{'output/prev:output'}(3),F|'prev_out_expand_dims'(1)], ctx=loop('time'[B])} (construction stack 'data_concat')>
prev_out1: <_TemplateLayer(CastLayer)(:template:cast) output/'prev_out1' out_type=Data{[B&Beam{'output/prev:output'}(3)], ctx=loop('time'[B])} (construction stack 'prev_out')>
prev_out0: <_TemplateLayer(ReinterpretDataLayer)(:template:reinterpret_data) output/'prev_out0' out_type=Data{[B&Beam{'output/prev:output'}(3)], dtype='int32', ctx=loop('time'[B])} (construction stack 'prev_out1')>

EXCEPTION
...
File "/mnt/projects/i6/returnn/returnn/tf/util/data.py", line 3256, in Data.placeholder
    line: self.sanity_check(assume_complete=False)
    locals:
      self = <local> Data{'concat_data_new_prev_out', [B&Beam{'output/prev:output'}(3),T|'new-time'[B],F|F'concat_data_new_prev_out_feature'(6)]}
      self.sanity_check = <local> <bound method Data.sanity_check of Data{'concat_data_new_prev_out', [B&Beam{'output/prev:output'}(3),T|'new-time'[B],F|F'concat_data_new_prev_out_feature'(6)]}>
      assume_complete = <not found>
  File "/mnt/projects/i6/returnn/returnn/tf/util/data.py", line 1820, in Data.sanity_check
    line: assert tag.batch == self.batch or self.batch.is_broadcast()
    locals:
      tag = <local> DimensionTag{'new-time'[B]}
      tag.batch = <local> BatchInfo{B}
      self = <local> Data{'concat_data_new_prev_out', [B&Beam{'output/prev:output'}(3),T|'new-time'[B],F|F'concat_data_new_prev_out_feature'(6)]}
      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 whole stack trace is in: https://gist.github.com/robin-p-schmitt/d59570c322891fcbb49bc0a7d6944562.

The error does not occur when `data_new` uses the `time_tag` and it also does not occur during training.

I debugged the run for both the case when the new `DimensionTag` is used and the case when the old one is used to see where they differ, which is here: https://github.com/rwth-i6/returnn/blob/7efd41f470c74fe70fc75bec1383bca6da81fbc1/returnn/tf/layers/basic.py#L118.
More specifically, the crucial difference happens here: https://github.com/rwth-i6/returnn/blob/7efd41f470c74fe70fc75bec1383bca6da81fbc1/returnn/tf/util/data.py#L3837.
If the old `DimensionTag` (`time_tag`) is used, the beam information of the tag is set via https://github.com/rwth-i6/returnn/blob/7efd41f470c74fe70fc75bec1383bca6da81fbc1/returnn/tf/util/data.py#L3853. 
On the other hand, if the new `DimensionTag` (`new_time_tag`) is used, this line is not executed and the beam information is not added to the tag, which leads to the error.


Discussion:

Resolving patch:
diff --git a/returnn/tf/util/data.py b/returnn/tf/util/data.py
--- a/returnn/tf/util/data.py
+++ b/returnn/tf/util/data.py
@@ -3292,6 +3292,7 @@ def beam(self, beam):
     self._beam = beam
     if self._batch:
       self._batch = self._batch.copy_set_beam(beam=beam)
+      self._adapt_batch_consistent_dim_tags()
 
   @property
   def dim(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
@@ -382,6 +382,40 @@ def test_concat_sources_dim1():
     session.run(out.output.placeholder, feed_dict=feed_dict)
 
 
+def test_concat_new_dim_tag():
+  from returnn.tf.util.data import DimensionTag
+  with make_scope():
+    n_out = 5
+    time_tag = DimensionTag(DimensionTag.Types.Spatial, "time")
+    new_time_tag = DimensionTag(DimensionTag.Types.Spatial, "new-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_new": {
+        "class": "reinterpret_data", "from": "data",
+        "set_dim_tags": {"t": new_time_tag}},
+      "output": {"class": "rec", "from": "data", "unit": {
+        "prev_out0": {
+          "class": "reinterpret_data", "from": "prev:output", "set_sparse": False},
+        "prev_out1": {"class": "cast", "from": "prev_out0", "dtype": "float32"},
+        "prev_out": {"class": "expand_dims", "from": "prev_out1", "axis": "f"},
+        "data_concat": {
+          "class": "copy", "from": ["base:data_new", "prev_out"]
+        },
+        "data_red": {"class": "reduce", "from": "data_concat", "axis": "stag:new-time", "mode": "max"},
+        "output_prob": {"class": "softmax", "from": "data_red", "target": "classes", "loss": "ce"},
+        "output": {
+          "class": "choice", "from": "output_prob", "beam_size": 3, "target": "classes",
+          "input_type": "prob", "initial_output": 0}
+      }}
+    })
+
+
 def test_ConcatLayer():
   with make_scope() as session:
     net_dict = {