Repository: rwth-i6/returnn

Issue:
ReinterpretDataLayer: set_dim_tags with 2D dynamic size
Using the following network:

```
network = {
      "start": {"class": "reinterpret_data", "from": "data:classes", "set_sparse": False},
      "slices": {"class": "slice_nd", "from": "data", "start": "start", "size": None},
      "output": {
        "class": "reinterpret_data", "from": "slices", "set_dim_tags": {"stag:sliced-time:slices": new_slice_tag}}
    }
```

with `new_slice_tag = DimensionTag(kind=DimensionTag.Types.Spatial, description="new-slice")`

leads to the following error:

```
layer <network via test_SliceNdLayer_ReinterpretDataLayer>/'data' output: Data{'data', [B,T|'time:var:extern_data:data'[B],F|F'feature:data'(5)]}
layer <network via test_SliceNdLayer_ReinterpretDataLayer>/'data:classes' output: Data{'classes', [B,T|'time:var:extern_data:classes'[B]], dtype='int32', sparse_dim=DimensionTag{F'classes:sparse-dim'(5)}, available_for_inference=False}
layer <network via test_SliceNdLayer_ReinterpretDataLayer>/'start' output: Data{'start_output', [B,T|'time:var:extern_data:classes'[B]], dtype='int32', available_for_inference=False}
layer <network via test_SliceNdLayer_ReinterpretDataLayer>/'slices' output: Data{'slices_gather_output', [B,T|'time:var:extern_data:classes'[B],'sliced-time:slices'[?],F|F'feature:data'(5)], available_for_inference=False}
layer <network via test_SliceNdLayer_ReinterpretDataLayer>/'output' output: Data{'output_output', [B,T|'time:var:extern_data:classes'[B],'new-slice'[?],F|F'feature:data'(5)], available_for_inference=False}
Exception creating layer <network via test_SliceNdLayer_ReinterpretDataLayer>/'output' of class ReinterpretDataLayer with opts:
{'_name': 'output',
 '_network': <TFNetwork '<network via test_SliceNdLayer_ReinterpretDataLayer>' train>,
 'name': 'output',
 'network': <TFNetwork '<network via test_SliceNdLayer_ReinterpretDataLayer>' train>,
 'output': Data{'output_output', [B,T|'time:var:extern_data:classes'[B],'new-slice'[?],F|F'feature:data'(5)], available_for_inference=False},
 'set_dim_tags': {'stag:sliced-time:slices': DimensionTag{'new-slice'[?]}},
 'sources': [<SliceNdLayer 'slices' out_type=Data{[B,T|'time:var:extern_data:classes'[B],'sliced-time:slices'[B,'time:var:extern_data:classes'[B]],F|F'feature:data'(5)], available_for_inference=False}>]}
EXCEPTION
...
File "/mnt/projects/i6/returnn/returnn/tf/util/data.py", line 470, in DimensionTag.set_tag_on_size_tensor
    line: self.dyn_size = x
    locals:
      self = <local> DimensionTag{'new-slice'[?]}
      self.dyn_size = <local> None
      x = <local> <tf.Tensor 'output/new-slice_size:0' shape=(?, ?) dtype=int32>
  File "/mnt/projects/i6/returnn/returnn/tf/util/data.py", line 344, in DimensionTag.dyn_size
    line: assert isinstance(dyn_size, tf.Tensor) and dyn_size.shape.ndims == 1
    locals:
      isinstance = <builtin> <built-in function isinstance>
      dyn_size = <local> <tf.Tensor 'output/new-slice_size:0' shape=(?, ?) dtype=int32>
      tf = <global> <module 'tensorflow' from '/home/robin/miniconda3/envs/i6_env/lib/python3.9/site-packages/tensorflow/__init__.py'>
      tf.Tensor = <global> <class 'tensorflow.python.framework.ops.Tensor'>
      dyn_size.shape = <local> TensorShape([Dimension(None), Dimension(None)]), len = 2
      dyn_size.shape.ndims = <local> 2
AssertionError
```

The full traceback is in https://gist.github.com/robin-p-schmitt/b9d3b266835cd0f869beb9201dbca0d8



Discussion:
Yea, `self.dyn_size = x` is not possible for non-standard dyn sizes (anything which is not shape [B]). It needs to assign `dyn_size_ext`.
> Yea, `self.dyn_size = x` is not possible for non-standard dyn sizes (anything which is not shape [B]). It needs to assign `dyn_size_ext`.

https://github.com/rwth-i6/returnn/blob/b6d0c47e5d445e419fa3eee2d295c16a8b755854/returnn/tf/layers/basic.py#L4213
https://github.com/rwth-i6/returnn/blob/b6d0c47e5d445e419fa3eee2d295c16a8b755854/returnn/tf/layers/basic.py#L4214

Swapping the lines like:

```
new_tag.dyn_size_ext = new_dyn_size_ext
new_tag.set_tag_on_size_tensor(new_dyn_size_ext.placeholder)
```

solves the error. Should I change that?

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
@@ -4145,8 +4145,8 @@ def __init__(self, switch_axes=None, size_base=None, set_axes=None,
           # Need to create new size tensor as long as we have get_tag_from_size_tensor.
           new_dyn_size_ext.placeholder = tf.identity(
             new_dyn_size_ext.placeholder, name=get_valid_scope_name_from_str(new_dyn_size_ext.name))
-          new_tag.set_tag_on_size_tensor(new_dyn_size_ext.placeholder)
           new_tag.dyn_size_ext = new_dyn_size_ext
+          new_tag.set_tag_on_size_tensor(new_dyn_size_ext.placeholder)
     self.output.placeholder = input_data.placeholder
     if len(self.sources) == 1:
       self.output_loss = self.sources[0].output_loss


Test patch:
diff --git a/tests/test_TFNetworkLayer.py b/tests/test_TFNetworkLayer.py
--- a/tests/test_TFNetworkLayer.py
+++ b/tests/test_TFNetworkLayer.py
@@ -3511,6 +3511,29 @@ def test_SliceNdLayer_set_tag_on_size_tensor():
     })
 
 
+def test_SliceNdLayer_ReinterpretDataLayer():
+  """
+  https://github.com/rwth-i6/returnn/issues/851
+  """
+  from returnn.tf.util.data import DimensionTag
+  new_slice_tag = DimensionTag(kind=DimensionTag.Types.Spatial, description="new-slice")
+  with make_scope():
+    n_out = 5
+    config = Config({
+      "debug_print_layer_output_template": True,
+      "extern_data": {
+        "data": {"dim": n_out},
+        "classes": {"dim": n_out, "sparse": True}
+      }})
+    net = TFNetwork(config=config, train_flag=True)
+    net.construct_from_dict({
+      "start": {"class": "reinterpret_data", "from": "data:classes", "set_sparse": False},
+      "slices": {"class": "slice_nd", "from": "data", "start": "start", "size": None},
+      "output": {
+        "class": "reinterpret_data", "from": "slices", "set_dim_tags": {"stag:sliced-time:slices": new_slice_tag}}
+    })
+
+
 def test_WindowLayer_output_placeholder():
   with make_scope() as session:
     net = TFNetwork(extern_data=ExternData())