Repository: rwth-i6/returnn

Issue:
Merge batch with modified time dim
In a case where the time dim is modified (e.g. due to downsampling/padding in a convolution) and we want to merge it with the batch dim, we currently face the problem that the creation of a `PaddedDim` in the `MergeDimsLayer` does not work because it is not possible to get the dim value for the modified time dim. This became obvious in https://github.com/rwth-i6/pytorch-to-returnn/issues/96

A demo test case which reproduces the issue looks like this and can also be found in #916:
```
def test_MergeDimsLayer_modified_time_dim():
  n_batch, n_time, n_in = 3, 7, 2
  config = Config({
    "extern_data": {"data": {"dim": n_in}},
    "debug_print_layer_output_template": True,
    "behavior_version": 12,
  })
  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict({
      "conv": {
        "class": "conv", "n_out": n_in, "filter_size": (3,), "strides": (2,), "padding": "valid", "from": "data:data"},
      "output": {"class": "merge_dims", "from": "conv", "axes": ["B", "T"], "keep_order": True},
    })
    in_data = net.extern_data.data["data"]
    out_t = net.get_default_output_layer().output.placeholder
    in_v = numpy.arange(0, n_batch * n_time * n_in).reshape((n_time, n_batch, n_in)).transpose(1, 0, 2)
    in_seq_lens = [n_time] * n_batch
    out_v = session.run(out_t, feed_dict={
      in_data.placeholder: in_v,
      in_data.size_placeholder[0]: in_seq_lens})
    assert isinstance(out_v, numpy.ndarray)
    numpy.testing.assert_equal(out_v, in_v)
```

The stack trace looks like this ([via](https://github.com/rwth-i6/returnn/runs/4967644403)):
```
ERROR: test_TFNetworkLayer.test_MergeDimsLayer_modified_time_dim
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/home/runner/.local/lib/python3.7/site-packages/nose/case.py", line 198, in TestBase.runTest
    line: self.test(*self.arg)
    locals:
      self = <local> test_TFNetworkLayer.test_MergeDimsLayer_modified_time_dim
      self.test = <local> <function test_MergeDimsLayer_modified_time_dim at 0x7f69d4a7def0>
      self.arg = <local> ()
  File "/home/runner/work/returnn/returnn/tests/test_TFNetworkLayer.py", line 2516, in test_MergeDimsLayer_modified_time_dim
    line: "output": {"class": "merge_dims", "from": "conv", "axes": ["B", "T"], "keep_order": True},
    locals:
       no locals
  File "/home/runner/work/returnn/returnn/returnn/tf/network.py", line 615, in TFNetwork.construct_from_dict
    line: self.construct_layer(net_dict, name, get_layer=get_layer)
    locals:
      self = <local> <TFNetwork '<network via test_MergeDimsLayer_modified_time_dim>' train=False>
      self.construct_layer = <local> <bound method TFNetwork.construct_layer of <TFNetwork '<network via test_MergeDimsLayer_modified_time_dim>' train=False>>
      net_dict = <local> {'conv': {'class': 'conv', 'n_out': 2, 'filter_size': (3,), 'strides': (2,), 'padding': 'valid', 'from': 'data:data'}, 'output': {'class': 'merge_dims', 'from': 'conv', 'axes': ['B', 'T'], 'keep_order': True}}
      name = <local> 'output', len = 6
      get_layer = <local> None
  File "/home/runner/work/returnn/returnn/returnn/tf/network.py", line 942, in TFNetwork.construct_layer
    line: return add_layer(name=name_with_prefix, layer_class=layer_class, **layer_desc)
    locals:
      add_layer = <local> <bound method TFNetwork.add_layer of <TFNetwork '<network via test_MergeDimsLayer_modified_time_dim>' train=False>>
      name = <local> 'output', len = 6
      name_with_prefix = <local> 'output', len = 6
      layer_class = <local> <class 'returnn.tf.layers.basic.MergeDimsLayer'>
      layer_desc = <local> {'axes': ['B', 'T'], 'keep_order': True, '_network': <TFNetwork '<network via test_MergeDimsLayer_modified_time_dim>' train=False>, '_name': 'output', 'sources': [<ConvLayer 'conv' out_type=Data{[B,T|'conv:conv:s0'[B],F|F'conv:channel'(2)]}>]}
  File "/home/runner/work/returnn/returnn/returnn/tf/network.py", line 1089, in TFNetwork.add_layer
    line: layer = self._create_layer(name=name, layer_class=layer_class, **layer_desc)
    locals:
      layer = <not found>
      self = <local> <TFNetwork '<network via test_MergeDimsLayer_modified_time_dim>' train=False>
      self._create_layer = <local> <bound method TFNetwork._create_layer of <TFNetwork '<network via test_MergeDimsLayer_modified_time_dim>' train=False>>
      name = <local> 'output', len = 6
      layer_class = <local> <class 'returnn.tf.layers.basic.MergeDimsLayer'>
      layer_desc = <local> {'axes': ['B', 'T'], 'keep_order': True, '_network': <TFNetwork '<network via test_MergeDimsLayer_modified_time_dim>' train=False>, '_name': 'output', 'sources': [<ConvLayer 'conv' out_type=Data{[B,T|'conv:conv:s0'[B],F|F'conv:channel'(2)]}>]}
  File "/home/runner/work/returnn/returnn/returnn/tf/network.py", line 991, in TFNetwork._create_layer
    line: layer_desc["output"] = layer_class.get_out_data_from_opts(**layer_desc)
    locals:
      layer_desc = <local> {'axes': ['B', 'T'], 'keep_order': True, '_network': <TFNetwork '<network via test_MergeDimsLayer_modified_time_dim>' train=False>, '_name': 'output', 'sources': [<ConvLayer 'conv' out_type=Data{[B,T|'conv:conv:s0'[B],F|F'conv:channel'(2)]}>], 'name': 'output', 'network': <TFNetwork '<network via..., len = 7
      layer_class = <local> <class 'returnn.tf.layers.basic.MergeDimsLayer'>
      layer_class.get_out_data_from_opts = <local> <bound method MergeDimsLayer.get_out_data_from_opts of <class 'returnn.tf.layers.basic.MergeDimsLayer'>>
  File "/home/runner/work/returnn/returnn/returnn/tf/layers/basic.py", line 3212, in MergeDimsLayer.get_out_data_from_opts
    line: data.batch = data.batch.copy_extend_with_padded_or_fixed_dim_tag(
            dim_tag=input_data.get_dim_tag(axis),
            batch_major=(axis > input_data.batch_dim_axis) if keep_order else True)
    locals:
      data = <local> Data{'output_output', [B,F|F'conv:channel'(2)]}
      data.batch = <local> BatchInfo{B}
      data.batch.copy_extend_with_padded_or_fixed_dim_tag = <local> <bound method BatchInfo.copy_extend_with_padded_or_fixed_dim_tag of BatchInfo{B}>
      dim_tag = <not found>
      input_data = <local> Data{'conv_output', [B,T|'conv:conv:s0'[B],F|F'conv:channel'(2)]}
      input_data.get_dim_tag = <local> <bound method Data.get_dim_tag of Data{'conv_output', [B,T|'conv:conv:s0'[B],F|F'conv:channel'(2)]}>
      axis = <local> 1
      batch_major = <not found>
      input_data.batch_dim_axis = <local> 0
      keep_order = <local> True
  File "/home/runner/work/returnn/returnn/returnn/tf/util/data.py", line 2278, in BatchInfo.copy_extend_with_padded_or_fixed_dim_tag
    line: new_dim = self._make_padded_dim(dim_tag)
    locals:
      new_dim = <not found>
      self = <local> BatchInfo{B}
      self._make_padded_dim = <local> <bound method BatchInfo._make_padded_dim of BatchInfo{B}>
      dim_tag = <local> Dim{'conv:conv:s0'[B]}
  File "/home/runner/work/returnn/returnn/returnn/tf/util/data.py", line 2182, in BatchInfo._make_padded_dim
    line: new_dim = BatchInfo.PaddedDim(dim_tag=dim_tag_base)
    locals:
      new_dim = <not found>
      BatchInfo = <global> <class 'returnn.tf.util.data.BatchInfo'>
      BatchInfo.PaddedDim = <global> <class 'returnn.tf.util.data.BatchInfo.PaddedDim'>
      dim_tag = <local> Dim{'conv:conv:s0'[B]}
      dim_tag_base = <local> Dim{'conv:conv:s0'[?]}
  File "/home/runner/work/returnn/returnn/returnn/tf/util/data.py", line 1860, in BatchInfo.PaddedDim.__init__
    line: super(BatchInfo.PaddedDim, self).__init__(size=dim_tag.get_dim_value())
    locals:
      super = <builtin> <class 'super'>
      BatchInfo = <global> <class 'returnn.tf.util.data.BatchInfo'>
      BatchInfo.PaddedDim = <global> <class 'returnn.tf.util.data.BatchInfo.PaddedDim'>
      self = <local> !AttributeError: 'PaddedDim' object has no attribute 'dim_tag'
      __init__ = <not found>
      size = <not found>
      dim_tag = <local> Dim{'conv:conv:s0'[?]}
      dim_tag.get_dim_value = <local> <bound method Dim.get_dim_value of Dim{'conv:conv:s0'[?]}>
  File "/home/runner/work/returnn/returnn/returnn/tf/util/data.py", line 954, in Dim.get_dim_value
    line: raise Exception('%s: need placeholder, self.dimension or self.dyn_size for dim value' % self)
    locals:
      Exception = <builtin> <class 'Exception'>
      self = <local> Dim{'conv:conv:s0'[?]}
Exception: Dim{'conv:conv:s0'[?]}: need placeholder, self.dimension or self.dyn_size for dim value
```


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
@@ -2176,7 +2176,7 @@ def _make_padded_dim(self, dim_tag):
     assert self.virtual_dims
     root = self.get_global_base()
     assert dim_tag.dyn_size is not None
-    dim_tag_base = dim_tag.get_same_base()
+    dim_tag_base = dim_tag.get_for_batch_ctx(self, dim_tag.control_flow_ctx)
     if dim_tag_base in root._global_padded_dims_by_dim_tag:
       return root._global_padded_dims_by_dim_tag[dim_tag_base]
     new_dim = BatchInfo.PaddedDim(dim_tag=dim_tag_base)


Test patch:
diff --git a/tests/test_TFNetworkLayer.py b/tests/test_TFNetworkLayer.py
--- a/tests/test_TFNetworkLayer.py
+++ b/tests/test_TFNetworkLayer.py
@@ -2501,6 +2501,25 @@ def test_MergeDimsLayer_2d_dynamic_merge_axis():
     assert numpy.all(sliced_time * feat_dim.dimension == merged_time)
 
 
+def test_MergeDimsLayer_modified_time_dim():
+  n_batch, n_time, n_in = 3, 7, 2
+  config = Config({
+    "extern_data": {"data": {"dim": n_in}},
+    "debug_print_layer_output_template": True,
+    "behavior_version": 12,
+  })
+  with make_scope() as session:
+    net = TFNetwork(config=config)
+    net.construct_from_dict({
+      "conv": {
+        "class": "conv", "n_out": n_in, "filter_size": (3,), "strides": (2,), "padding": "valid", "from": "data:data"},
+      "output": {"class": "merge_dims", "from": "conv", "axes": ["B", "T"], "keep_order": True},
+    })
+    out = net.get_default_output_layer().output
+    net.initialize_params(session)
+    session.run(out.placeholder, feed_dict=make_feed_dict(net.extern_data))
+
+
 def test_FlattenBatchLayer():
   from returnn.tf.util.data import BatchInfo
   n_batch, n_time, n_in = 3, 4, 2