Repository: rwth-i6/returnn

Issue:
Dim tags matching should use `allow_same_spatial_dim=False`
This is kind of an follow up to #666.

I have the problem currently, that `SpatialDim("a", 1)` matches `SpatialDim("b", 1)` in many cases, e.g. `get_common_data`, `copy_compatible_to` and `find_matching_dim_map` used by `GatherLayer`/`DotLayer`/etc.

Looking at the `Dim.is_equal` code, `allow_same_spatial_dim` matches exactly if
1. both tags are spatial dims (or feature if `treat_feature_as_spatial` which e.g. `get_common_data` sets)
2. the dimensions of `self` and `other` are equal, but not `None`

For my case, I definitely do not want these axes to match (they are not related in any way).

But this kind of matching is currently necessary for something like this:
```
      "lin1": {"class": "linear", "activation": "sigmoid", "n_out": 5, "from": "data:data"},
      "lin2": {"class": "linear", "activation": "sigmoid", "n_out": 5, "from": "data:data"},
      "output": {"class": "combine", "mode": "add", "from": ["lin1", "lin2"]},  # would give [B,T,F1,F2] with allow_same_spatial_dim=False!!
```

This is also because currently, `LinearLayer` does not set `derived_from_tag` on its output feature dim.
Then, `derived_matches` would catch this.
We can fix `LinearLayer` and to do this properly (currently this is via `_base_get_out_data_from_opts`, when creating a feature dim tag, I think it should just derive it from its input feature dim tag).
For user specified `out_dim`s however, we would not automatically set derived from to the input dim tag.

For other cases, this is more difficult maybe.
E.g. if you have two inputs with different time axes, and then split into the same dim size.
Then, the user would maybe want these axes to match, but they are in no way derived from the same thing.

Another solution is, similar to what @albertz wrote in #666, to separate user defined dim tags and automatically created ones. Then, user defined ones could have a more restrictive matching logic.


Discussion:
`derived_from_tag`/`derived_matches` is not the right option here. `derived_from_tag` means that it is derived from this tag (via some op, like +- or so). We use it as a heuristic to assume in the rec layer, when some dim is derived (but unknown yet due to template construction, also this was before the dim math), that it is the same as the orig dim. E.g. we had the case where some people did some extra padding, then convolution with `padding="valid"` which removed the padding again such that it became the same, and that inside a rec layer, on the attention encoder axis.

The output feature dim is usually never the same as the input feature dim. And also, it is not derived from the input feature dim in any way. They are completely independent. So `derived_from_tag` doesn't really make sense here.
Ah yeah, you're right, the `derived_*` logic is not applicable here.

So really, I think specifying `n_out: int` is the problem at the moment. In many cases currently, if `n_in == n_out`, we want the in and output dim tags to match. Like in my example above.
Perhaps we could change it so that if `n_out: int` is specified, and `n_out == in_dim.dimension`, that then the newly created output dim tag will match with the input dim tag (either by being exactly the same, or via some other logic).
I wonder if this most cases where we needed `allow_same_spatial_dim` before.

Logically, this does not make sense 100%, because the in and out dims should logically be different, but we matched them before always anyway, and in the future, we discourage anyone from using `n_out: int` anyway.
Also the  current matching logic is just dangerously broken and depends on the order of axes, e.g. this fails currently:
```python
def test_same_spatial_dim():
  foo_dim = SpatialDim("foo", 3)
  bar_dim = SpatialDim("bar", 3)
  x = Data("a", dim_tags=[foo_dim, bar_dim])
  y = Data("b", dim_tags=[bar_dim, foo_dim])
  # test with default is_equal_opts=None
  assert_equal(x.find_matching_dim_map(y, other_axes=[0, 1]), {0: 1, 1: 0})  # returns {0: 0, 1: 1} currently
```
We should probably just at least throw an error here.
Note that I stumbled upon this (or related) specifically for the `DotLayer` in #1154, and fixed via #1155.
In #1155, I introduced a new behavior version to trigger the more strict matching for `DotLayer`. This is anyway only for the case with `_auto_var_axes`, i.e. when `var1` and `var2` are not specified explicitly.

But actually, I did not fully think through it, whether this is maybe difficult for users who did not adapt to use dim tags. Maybe not for `DotLayer`, as you can always specify everything explicitly.
I'm now thinking whether my proposed solution from from #666 is maybe really the better way which would have avoided these issues:

> As another solution, maybe we can disallow the broadcasting (or other `is_equal` attribs) if this is an explicit dim tag by the user? To implement that, we would need to go through all places in RETURNN where a static dim tag is created, and add a new flag like `auto_created=True` or so. And then allow the broadcasting only if `auto_created and dimension == 1`.

Again we should think about the quality (#634). When comparing a user-generated dim tag to another user-generated dim tag, it should be strict. When comparing an auto-created dim tag to another auto-created dim tag, it can use the `is_equal` attribs and be more relaxed. When comparing an auto-created dim tag to a user-created dim tag, what then? I think it should also be strict.

But yes, we also should think about having it all really well defined and the code should be clean. See #975.

We have the `auto_generated` flag.

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
@@ -7904,7 +7904,9 @@ def __init__(self, mask, unit, masked_from,
           new_size = masked_from.output.get_sequence_lengths()
         else:
           new_size = idxs[-1]  # [B]
-          out_spatial_dim.get_for_batch_ctx(self.output.batch, self.output.control_flow_ctx).dyn_size = new_size
+          out_spatial_dim = out_spatial_dim.get_for_batch_ctx(self.output.batch, self.output.control_flow_ctx)
+          if out_spatial_dim.dyn_size is None:
+            out_spatial_dim.dyn_size = new_size
         new_time = tf.reduce_max(new_size)  # T'
         idxs = where_bc(mask_t, idxs - 1, new_time)
 
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
@@ -126,6 +126,8 @@ def __init__(self, kind=Types.Unspecified, description=None,
     self._undefined = undefined
     self.generic = generic
     self.special = special
+    if derived_from_tag:
+      auto_generated = derived_from_tag.auto_generated
     self.auto_generated = auto_generated
     # We can have different tag variants per batch info (e.g. with beam), or per control flow ctx.
     # They each have same_as = self. The same_base should have the base (global) batch info.
@@ -760,6 +762,7 @@ def is_equal(self, other, ignore_feature_dim=False, allow_same_feature_dim=False
     :param bool derived_matches:
     :rtype: bool
     """
+    from returnn.util import BehaviorVersion
     if self is other:  # first some fast path check
       return True
     if self.special or other.special:
@@ -806,6 +809,16 @@ def is_equal(self, other, ignore_feature_dim=False, allow_same_feature_dim=False
       # or when we used MergeDimsLayer on the batch axis, or so.
       # We might need to extend the logic here later.
       return True
+    if BehaviorVersion.get() >= 16:
+      # Either self or other is some dim tag explicitly created by the user,
+      # and they are not the same, so we never treat them as equal.
+      if not self.auto_generated or not other.auto_generated:
+        if broadcast_matches and (
+              (self.dimension == 1 and self.auto_generated) or
+              (other.dimension == 1 and other.auto_generated)):
+          pass  # exception, allow broadcast logic
+        else:
+          return False
     if self_kind == other_kind == self.Types.Feature:
       if allow_same_feature_dim:
         return True
@@ -1059,6 +1072,7 @@ def declare_same_as(self, other):
       other_same_base._vocab = self._vocab
     elif other_same_base._vocab and not self._vocab:
       self._vocab = other_same_base._vocab
+    self.auto_generated = self_same_as.auto_generated = other_same_base.auto_generated
     # Take over derived_from_op. However, only if this would not introduce cycles!
     if not self_derived_bases.issuperset(other_derived_bases):
       if self.derived_from_op and not other_same_base.derived_from_op:
@@ -5608,8 +5622,8 @@ def map_other_axis_to_self(other_axis, taken_self_axes):
             other_axes, other, self, other_axis, opt)
         if matching:
           break
-      assert matching, 'cannot match the axes %s from %s to %s. Failing at axis %s' % (
-        other_axes, other, self, other_axis)
+      assert matching, 'cannot match the axes %s from %s to %s. Failing at axis %s, tag %s' % (
+        other_axes, other, self, other_axis, other.dim_tags[other_axis])
       # If there are multiple matches (e.g. because two axes have the same feature dim), leave their order intact.
       # We do this by always choosing the first unused match which is the smallest axes
       return matching[0]
@@ -5859,7 +5873,8 @@ def _infer_dim_tags_tuple_from_shape(
           description="%s:var:extern_data:%s" % (tag_name, name),
           # Spatial dim tag, even if axis == feature_dim_axis. This is to keep the old behavior.
           # This is such that Dim.is_equal behaves as before, e.g. in Data.get_common_data.
-          kind=Dim.Types.Spatial)
+          kind=Dim.Types.Spatial,
+          auto_generated=True)
         dim_tags[axis] = tag
       dyn_size = tag.dyn_size
     if tag:
diff --git a/returnn/util/basic.py b/returnn/util/basic.py
--- a/returnn/util/basic.py
+++ b/returnn/util/basic.py
@@ -215,7 +215,7 @@ class BehaviorVersion:
   The version will be set after the config is defined at __main__.init_config() or Engine.__init__()
   """
 
-  _latest_behavior_version = 15
+  _latest_behavior_version = 16
   _behavior_version = None  # type: typing.Optional[int]
 
   @classmethod


Test patch:
diff --git a/tests/test_TFNetworkLayer.py b/tests/test_TFNetworkLayer.py
--- a/tests/test_TFNetworkLayer.py
+++ b/tests/test_TFNetworkLayer.py
@@ -5031,7 +5031,7 @@ def test_ScatterNdLayer_RangeLayer():
           "out_type": {"shape": (), "dtype": "int32"}},  # (B,)
     "range": {"class": "range", "limit": n_ts, "out_spatial_dim": ts_dim},  # (Ts,)
     "add_t": {"class": "combine", "kind": "add", "from": ["t", "range"], "out_shape": {batch_dim, ts_dim}},  # (B,Ts)
-    "t_rel_var": {"class": "variable", "shape": (n_ts, n_out), "init": "glorot_uniform"},  # (B,Ts,D)
+    "t_rel_var": {"class": "variable", "shape": (ts_dim, n_out), "init": "glorot_uniform"},  # (B,Ts,D)
     "output": {"class": "scatter_nd", "from": "t_rel_var", "position": "add_t", "position_axis": ts_dim,
                "output_dim_via_time_from": "data", "filter_invalid_indices": True}
   }
@@ -6138,6 +6138,25 @@ def test_GatherLayer_broadcast_dim():
   })
 
 
+def test_GatherLayer_different_static_dims():
+  # https://github.com/rwth-i6/returnn/issues/1219
+  from returnn.tf.util.data import batch_dim
+  head_dim = SpatialDim("head", 2)
+  round_dim = SpatialDim("round", 2)
+  chunk_dim = SpatialDim("chunk")
+  time_dim = SpatialDim("time")
+  config = Config({"extern_data": {
+    "source": {"dim_tags": [batch_dim, head_dim, time_dim]},
+    "position": {"dim_tags": [batch_dim, round_dim, chunk_dim], "dtype": "int32"}},
+    "debug_print_layer_output_template": True})
+  net = TFNetwork(config=config)
+  net.construct_from_dict({
+    "output": {
+      'class': 'gather', 'from': 'data:source', 'position': 'data:position', 'axis': time_dim,
+      'out_shape': {batch_dim, head_dim, round_dim, chunk_dim}}
+  })
+
+
 def test_SliceNdLayer():
   n_batch = 5
   n_time = 7
diff --git a/tests/test_TFNetworkRecLayer.py b/tests/test_TFNetworkRecLayer.py
--- a/tests/test_TFNetworkRecLayer.py
+++ b/tests/test_TFNetworkRecLayer.py
@@ -5708,6 +5708,10 @@ def test_reclayer_att_weights_output_layer():
     },
   }})
 
+  att_heads = Dim(kind=Dim.Types.Spatial, description="att_heads", dimension=1)
+  att_t = Dim(kind=Dim.Types.Spatial, description="att_t")
+  label_axis = Dim(kind=Dim.Types.Spatial, description="label-axis")
+
   net_dict = {
     "encoder": {"class": "copy", "from": "data"},
     "existing_alignment": {"class": "copy", "from": "data:alignment"},
@@ -5762,9 +5766,7 @@ def test_reclayer_att_weights_output_layer():
           "from": "att_energy0",
           "is_output_layer": False,
           "set_dim_tags": {
-            "f": Dim(
-              kind=Dim.Types.Spatial, description="att_heads", dimension=1
-            )
+            "f": att_heads
           },
         },
         "att_energy0": {
@@ -5794,9 +5796,7 @@ def test_reclayer_att_weights_output_layer():
           "class": "reinterpret_data",
           "from": "att_val_split0",
           "set_dim_tags": {
-            "dim:1": Dim(
-              kind=Dim.Types.Spatial, description="att_heads", dimension=1
-            )
+            "dim:1": att_heads
           },
         },
         "att_val_split0": {
@@ -5858,9 +5858,7 @@ def test_reclayer_att_weights_output_layer():
           "class": "reinterpret_data",
           "from": "segments0",
           "set_dim_tags": {
-            "stag:sliced-time:segments": Dim(
-              kind=Dim.Types.Spatial, description="att_t"
-            )
+            "stag:sliced-time:segments": att_t
           },
         },
         "segments0": {
@@ -5931,7 +5929,7 @@ def test_reclayer_att_weights_output_layer():
       "class": "masked_computation",
       "from": "output/segment_lens",
       "mask": "is_label",
-      "out_spatial_dim": Dim(kind=Dim.Types.Spatial, description="label-axis"),
+      "out_spatial_dim": label_axis,
       "register_as_extern_data": "segment_lens_masked",
       "unit": {"class": "copy", "from": "data"},
     },
@@ -5939,7 +5937,7 @@ def test_reclayer_att_weights_output_layer():
       "class": "masked_computation",
       "from": "output/segment_starts",
       "mask": "is_label",
-      "out_spatial_dim": Dim(kind=Dim.Types.Spatial, description="label-axis"),
+      "out_spatial_dim": label_axis,
       "register_as_extern_data": "segment_starts_masked",
       "unit": {"class": "copy", "from": "data"},
     },
@@ -7015,8 +7013,8 @@ def test_GenericAttentionLayer_extra_spatial_multi_head():
   net = TFNetwork(extern_data=ExternData(), config=Config({"debug_print_layer_output_template": True}))
   dec_time = SpatialDim("dec time")
   enc_time = SpatialDim("enc time")
-  num_heads = 8
-  heads_dim = FeatureDim("heads", dimension=num_heads)
+  heads_dim = FeatureDim("heads", dimension=8)
+  feat_dim = FeatureDim("feat", dimension=2048)
   kwargs = dict(
     name="att", network=net,
     weights=InternalLayer(
@@ -7026,14 +7024,13 @@ def test_GenericAttentionLayer_extra_spatial_multi_head():
     base=InternalLayer(
       name="enc_value", network=net,
       output=Data(
-        name='enc_value_output', shape=(None, num_heads, 2048), batch_dim_axis=1, auto_create_placeholders=True,
-        same_dim_tags_as={"t": enc_time})))
+        name='enc_value_output', dim_tags=[enc_time, batch_dim, heads_dim, feat_dim], auto_create_placeholders=True)))
   print("GenericAttentionLayer kwargs:")
   pprint(kwargs)
   kwargs["output"] = GenericAttentionLayer.get_out_data_from_opts(**kwargs)
   layer = GenericAttentionLayer(**kwargs)
   layer.output.sanity_check()
-  assert layer.output.shape == (num_heads, None, 2048) and layer.output.have_time_axis()
+  assert layer.output.shape == (heads_dim.dimension, None, feat_dim.dimension) and layer.output.have_time_axis()
   assert len(layer.output.size_placeholder) == 1
   assert list(layer.output.size_placeholder.values())[0] is layer.weights.output.size_placeholder[0]
 
@@ -8479,7 +8476,7 @@ def combine_soft_hard_att(self, source, **kwargs):
       "unit": {
         "s_transformed": {"class": "linear", "activation": None, "with_bias": False, "from": ["s"],
                           "n_out": EncKeyTotalDim},
-        "t_rel_var": {"class": "variable", "shape": (6, EncKeyTotalDim)},
+        "t_rel_var": {"class": "variable", "shape": (t_rel_idxs_dim, EncKeyTotalDim)},
         "t_rel_idxs_": {"class": "range", "limit": 6, "out_spatial_dim": t_rel_idxs_dim},
         "prev_t_": {"class": "reinterpret_data", "set_sparse": False, "from": "prev:t"},
         "t_rel_idxs": {"class": "combine", "kind": "add", "from": ["prev_t_", "t_rel_idxs_"],