Repository: rwth-i6/returnn

Issue:
DotLayer heuristical dim matching causes problems
Specifically, these tests fail:
```python
def test_DotLayer_dim_wrong_matching_same_dim_value():
  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim
  time_dim = SpatialDim("time")
  feat_dim = FeatureDim("feat", dimension=5)
  feat2_dim = FeatureDim("other-feat", dimension=5)

  # First directly check DotLayer.get_out_data_from_opts.
  # This is more similar like we have it in returnn_common
  # and might trigger different errors due to the dim matching logic of DotLayer,
  # which behaves slightly different when there are no size_placeholders set yet,
  # see Dim.is_equal with unknown_spatial_matches.
  a = Data("a", dim_tags=[batch_dim, time_dim, feat_dim])
  b = Data("b", dim_tags=[batch_dim, time_dim, feat2_dim])
  net = TFNetwork(config=Config(), extern_data=ExternData())
  out = DotLayer.get_out_data_from_opts(
    name="dot",
    sources=[InternalLayer(name="a", network=net, output=a), InternalLayer(name="b", network=net, output=b)],
    reduce=time_dim)
  assert out.dim_tags == (batch_dim, feat_dim, feat2_dim)

  # Now full config.
  config = Config({
    "extern_data": {
      "a": {"dim_tags": [batch_dim, time_dim, feat_dim]},
      "b": {"dim_tags": [batch_dim, time_dim, feat2_dim]},
    },
    "network": {
      "output": {"class": "dot", "from": ["data:a", "data:b"], "reduce": time_dim},
    },
    "debug_runtime_sanity_checks": True,
  })
  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict(config.typed_dict["network"])
    layer = net.get_default_output_layer()
    assert layer.output.dim_tags == (batch_dim, feat_dim, feat2_dim)
    feed_dict = make_feed_dict(net.extern_data)
    session.run(layer.output.placeholder, feed_dict=feed_dict)


def test_DotLayer_dim_wrong_matching_derived():
  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim
  time_dim = SpatialDim("time")
  time_dim_2 = time_dim * 2
  assert time_dim_2.derived_from_tag == time_dim
  assert time_dim_2.get_same_derived_base() == time_dim
  feat_dim = FeatureDim("feat", dimension=5)

  # First directly check DotLayer.get_out_data_from_opts.
  # This is more similar like we have it in returnn_common
  # and might trigger different errors due to the dim matching logic of DotLayer,
  # which behaves slightly different when there are no size_placeholders set yet,
  # see Dim.is_equal with unknown_spatial_matches.
  a = Data("a", dim_tags=[batch_dim, time_dim, feat_dim])
  b = Data("b", dim_tags=[batch_dim, time_dim_2, feat_dim])
  net = TFNetwork(config=Config(), extern_data=ExternData())
  out = DotLayer.get_out_data_from_opts(
    name="dot",
    sources=[InternalLayer(name="a", network=net, output=a), InternalLayer(name="b", network=net, output=b)],
    reduce=feat_dim)
  assert out.dim_tags == (batch_dim, time_dim, time_dim_2)

  # Now full config.
  config = Config({
    "extern_data": {
      "a": {"dim_tags": [batch_dim, time_dim, feat_dim]},
      "b": {"dim_tags": [batch_dim, time_dim_2, feat_dim]},
    },
    "network": {
      "output": {"class": "dot", "from": ["data:a", "data:b"], "reduce": feat_dim},
    },
    "debug_runtime_sanity_checks": True,
  })
  with make_scope() as session:
    net = TFNetwork(config=config)
    net.construct_from_dict(config.typed_dict["network"])
    layer = net.get_default_output_layer()
    assert layer.output.dim_tags == (batch_dim, time_dim, time_dim_2)
    feed_dict = make_feed_dict(net.extern_data)
    session.run(layer.output.placeholder, feed_dict=feed_dict)
```
The first test fails with:
```
Traceback (most recent call last):
  File "/Users/az/i6/setups/2022-03-19--sis-i6-exp/ext/returnn/tests/test_TFNetworkLayer.py", line 7166, in test_DotLayer_dim_wrong_matching_same_dim_value
    line: assert out.dim_tags == (batch_dim, feat_dim, feat2_dim)
    locals:
      out = <local> Data{'dot_output', [B?,F|F'feat'(5)]}
      out.dim_tags = <local> (Dim{B}, Dim{F'feat'(5)})
      batch_dim = <local> Dim{B}
      feat_dim = <local> Dim{F'feat'(5)}
      feat2_dim = <local> Dim{F'other-feat'(5)}
AssertionError: assert (Dim{B}, Dim{F'feat'(5)}) == (Dim{B}, Dim{F'feat'(5)}, Dim{F'other-feat'(5)})
  Right contains one more item: Dim{F'other-feat'(5)}
  Full diff:
  - (Dim{B}, Dim{F'feat'(5)}, Dim{F'other-feat'(5)})
  + (Dim{B}, Dim{F'feat'(5)})
```

The second test fails with:
```
Traceback (most recent call last):
  File "/Users/az/i6/setups/2022-03-19--sis-i6-exp/ext/returnn/tests/test_TFNetworkLayer.py", line 7208, in test_DotLayer_dim_wrong_matching_derived
    line: assert out.dim_tags == (batch_dim, time_dim, time_dim_2)
    locals:
      out = <local> Data{'dot_output', [B?,T|'time'[?]]}
      out.dim_tags = <local> (Dim{B}, Dim{'time'[?]})
      batch_dim = <local> Dim{B}
      time_dim = <local> Dim{'time'[?]}
      time_dim_2 = <local> Dim{'time*2'[?]}
AssertionError: assert (Dim{B}, Dim{'time'[?]}) == (Dim{B}, Dim{'time'[?]}, Dim{'time*2'[?]})
  Right contains one more item: Dim{'time*2'[?]}
  Full diff:
  - (Dim{B}, Dim{'time'[?]}, Dim{'time*2'[?]})
  + (Dim{B}, Dim{'time'[?]})
```
(Actually it already fails with the `time_dim_2.derived_from_tag == time_dim` but this is an independent fix.)

The problematic code is in `DotLayer._auto_var_axes`:
```python
    is_equal_opts = dict(
      treat_feature_as_spatial=True, allow_same_spatial_dim=True,
      undefined_matches=True, derived_matches=True)
```
I think we have `derived_matches` mostly for `RecLayer` and its subnet template construction in some cases. But I don't really remember exactly anymore which cases. The `allow_same_spatial_dim` is for all networks where dim tags are not consistently used everywhere.

Note that this logic is only relevant when the user does not explicitly specify `var1` and `var2`, which is maybe not so common.

Anyway, a change of this behavior probably requires a new behavior version.


Discussion:

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
@@ -7137,11 +7137,14 @@ def _auto_var_axes(source1, source2, red1, red2):
     :rtype: (list[int], list[int])
     """
     from returnn.util import BehaviorVersion
-    is_equal_opts = dict(
-      treat_feature_as_spatial=True, allow_same_spatial_dim=True,
-      undefined_matches=True, derived_matches=True)
-    if BehaviorVersion.get() < 11:
-      is_equal_opts["broadcast_matches"] = True
+    if BehaviorVersion.get() >= 14:
+      is_equal_opts = {}
+    else:
+      is_equal_opts = dict(
+        treat_feature_as_spatial=True, allow_same_spatial_dim=True,
+        undefined_matches=True, derived_matches=True)
+      if BehaviorVersion.get() < 11:
+        is_equal_opts["broadcast_matches"] = True
     tags1 = [None if i in red1 else tag for (i, tag) in enumerate(source1.dim_tags)]
     tags2 = [None if i in red2 else tag for (i, tag) in enumerate(source2.dim_tags)]
     var1 = [
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
@@ -1666,7 +1666,26 @@ def as_dim(self):
       return Dim(
         kind=dim_kind, description="*".join(map(Dim._get_description, self.terms)),
         dimension=self.dimension,
-        derived_from_op=Dim.Op(kind="mul", inputs=list(self.terms)))
+        derived_from_op=Dim.Op(kind="mul", inputs=list(self.terms)),
+        derived_from_tag=self.representative_tag())
+
+    def representative_tag(self):
+      """
+      :rtype: Dim|None
+      """
+      # Also see Dim._OpLinearTerm.representative_tag().
+      # First find any dynamic.
+      for term_ in self.terms:
+        if term_.is_dynamic():
+          return term_
+      # Now find non-unspecified.
+      for term_ in self.terms:
+        if term_.kind != Dim.Types.Unspecified:
+          return term_
+      # Now find any.
+      for term_ in self.terms:
+        return term_
+      return None
 
   @classmethod
   def _get_description(cls, dim, brackets=True):
@@ -1844,7 +1863,7 @@ def representative_tag(self):
       # First find any dynamic.
       for term in self.terms:
         for term_ in term.terms:
-          if term_.dimension is None:
+          if term_.is_dynamic():
             return term_
       # Now find non-unspecified.
       for term in self.terms:
diff --git a/returnn/util/basic.py b/returnn/util/basic.py
--- a/returnn/util/basic.py
+++ b/returnn/util/basic.py
@@ -238,7 +238,7 @@ class BehaviorVersion:
   The version will be set after the config is defined at __main__.init_config() or Engine.__init__()
   """
 
-  _latest_behavior_version = 13
+  _latest_behavior_version = 14
   _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
@@ -7145,6 +7145,88 @@ def test_DotLayer_sparse_input():
     session.run(layer.output.placeholder, feed_dict=feed_dict)
 
 
+def test_DotLayer_dim_wrong_matching_same_dim_value():
+  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim
+  time_dim = SpatialDim("time")
+  feat_dim = FeatureDim("feat", dimension=5)
+  feat2_dim = FeatureDim("other-feat", dimension=5)
+
+  # First directly check DotLayer.get_out_data_from_opts.
+  # This is more similar like we have it in returnn_common
+  # and might trigger different errors due to the dim matching logic of DotLayer,
+  # which behaves slightly different when there are no size_placeholders set yet,
+  # see Dim.is_equal with unknown_spatial_matches.
+  a = Data("a", dim_tags=[batch_dim, time_dim, feat_dim])
+  b = Data("b", dim_tags=[batch_dim, time_dim, feat2_dim])
+  net = TFNetwork(config=Config(), extern_data=ExternData())
+  out = DotLayer.get_out_data_from_opts(
+    name="dot",
+    sources=[InternalLayer(name="a", network=net, output=a), InternalLayer(name="b", network=net, output=b)],
+    reduce=time_dim)
+  assert out.dim_tags == (batch_dim, feat_dim, feat2_dim)
+
+  # Now full config.
+  config = Config({
+    "extern_data": {
+      "a": {"dim_tags": [batch_dim, time_dim, feat_dim]},
+      "b": {"dim_tags": [batch_dim, time_dim, feat2_dim]},
+    },
+    "network": {
+      "output": {"class": "dot", "from": ["data:a", "data:b"], "reduce": time_dim},
+    },
+    "debug_runtime_sanity_checks": True,
+  })
+  with make_scope() as session:
+    net = TFNetwork(config=config)
+    net.construct_from_dict(config.typed_dict["network"])
+    layer = net.get_default_output_layer()
+    assert layer.output.dim_tags == (batch_dim, feat_dim, feat2_dim)
+    feed_dict = make_feed_dict(net.extern_data)
+    session.run(layer.output.placeholder, feed_dict=feed_dict)
+
+
+def test_DotLayer_dim_wrong_matching_derived():
+  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim
+  time_dim = SpatialDim("time")
+  time_dim_2 = time_dim * 2
+  assert time_dim_2.derived_from_tag == time_dim
+  assert time_dim_2.get_same_derived_base() == time_dim
+  feat_dim = FeatureDim("feat", dimension=5)
+
+  # First directly check DotLayer.get_out_data_from_opts.
+  # This is more similar like we have it in returnn_common
+  # and might trigger different errors due to the dim matching logic of DotLayer,
+  # which behaves slightly different when there are no size_placeholders set yet,
+  # see Dim.is_equal with unknown_spatial_matches.
+  a = Data("a", dim_tags=[batch_dim, time_dim, feat_dim])
+  b = Data("b", dim_tags=[batch_dim, time_dim_2, feat_dim])
+  net = TFNetwork(config=Config(), extern_data=ExternData())
+  out = DotLayer.get_out_data_from_opts(
+    name="dot",
+    sources=[InternalLayer(name="a", network=net, output=a), InternalLayer(name="b", network=net, output=b)],
+    reduce=feat_dim)
+  assert out.dim_tags == (batch_dim, time_dim, time_dim_2)
+
+  # Now full config.
+  config = Config({
+    "extern_data": {
+      "a": {"dim_tags": [batch_dim, time_dim, feat_dim]},
+      "b": {"dim_tags": [batch_dim, time_dim_2, feat_dim]},
+    },
+    "network": {
+      "output": {"class": "dot", "from": ["data:a", "data:b"], "reduce": feat_dim},
+    },
+    "debug_runtime_sanity_checks": True,
+  })
+  with make_scope() as session:
+    net = TFNetwork(config=config)
+    net.construct_from_dict(config.typed_dict["network"])
+    layer = net.get_default_output_layer()
+    assert layer.output.dim_tags == (batch_dim, time_dim, time_dim_2)
+    feed_dict = make_feed_dict(net.extern_data)
+    session.run(layer.output.placeholder, feed_dict=feed_dict)
+
+
 def test_subnet_load_on_init():
   import tempfile
   model_tmp_dir = tempfile.mkdtemp("tmp-checkpoint")
diff --git a/tests/test_TFUtil.py b/tests/test_TFUtil.py
--- a/tests/test_TFUtil.py
+++ b/tests/test_TFUtil.py
@@ -1571,6 +1571,14 @@ def test_dim_math_feat_declare_same_as_circle():
   assert s == {feat_dim, feat2_dim}
 
 
+def test_dim_math_derived():
+  from returnn.tf.util.data import SpatialDim
+  time_dim = SpatialDim("time")
+  time_dim_2 = time_dim * 2
+  assert time_dim_2.derived_from_tag == time_dim
+  assert time_dim_2.get_same_derived_base() == time_dim
+
+
 def test_sequence_mask_len_via_loop():
   seq_len = tf.while_loop(
     cond=lambda x: tf.less(x[0], 2),