Repository: rwth-i6/returnn

Issue:
`DotLayer` could infer `var1`/`var2` automatically
This is more of a proposal for now. The idea is as follows:

For e.g. `GatherLayer` we automatically infer the shared axes between the positions and the params automatically using the dim tags. For me, this logic worked very well in the past.
In a same way, we could also infer the `var1`/`var2` axes as be the axes which are not present in both inputs for `DotLayer`.
This would also work well for when axis change depending on automatic rec layer optimization, and make a `common` argument that could be specified instead of `var1`/`var2` like discussed in #569 redundant.

I'm not 100% sure if this might be unexpected behavior though. E.g., to compute an dyadic product of a `[F]`-shaped tensor (i.e. compute `v * v^T`), the user could previously use `DotLayer` and set `var1=var2=F` as well as `red1=red2=()`, and then it would do exactly that, giving a `[F1,F2]` tensor.
Now with that change, you would first need to create a copy of the input, rename the axis to something different, and then use `DotLayer`.
But the old solution anyway had some issues: It would not be clear what the output dim tags (here `F1` and `F2`) would be. They should would be the same (I'm really not sure what our implementation currently does), but then you could not reference  the axes uniquely. That would be pretty bad.
With the new solution, we are very explicit what the output dim tags are.

Similarly, we can also get away with just a single `reduce` argument (#636).

_Originally posted by @Zettelkasten in https://github.com/rwth-i6/returnn/issues/627#issuecomment-913051726_


Discussion:
Wow, we were almost simultaneously. :) I will leave this open and close mine (#737).

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
@@ -6268,7 +6268,7 @@ def __init__(self,
       message="DotLayer: Axes must be referenced by tag or special specified, not by int.",
       version=3)
     BehaviorVersion.require(
-      condition=all(a is not NotSpecified for a in (red1, red2, var1, var2)),
+      condition=all(a is not NotSpecified for a in (red1, red2)),
       message="DotLayer: Axes must be specified explicitly. There is no default.",
       version=3)
     BehaviorVersion.require(
@@ -6292,6 +6292,9 @@ def __init__(self,
     a_reduce_axes = a_out.get_axes_from_description(red1, allow_int=axis_desc_allow_int)
     b_reduce_axes = b_out.get_axes_from_description(red2, allow_int=axis_desc_allow_int)
     assert a_reduce_axes and b_reduce_axes, "%s: sources %r, red1 %r, red2 %r" % (self, self.sources, red1, red2)
+    if BehaviorVersion.get() >= 3 and (var1 is NotSpecified or var2 is NotSpecified):
+      assert var1 is NotSpecified and var2 is NotSpecified
+      var1, var2 = self._auto_var_axes(a_out, b_out, a_reduce_axes, b_reduce_axes)
     a_var_axes = a_out.get_axes_from_description(var1, allow_int=axis_desc_allow_int)
     b_var_axes = b_out.get_axes_from_description(var2, allow_int=axis_desc_allow_int)
     assert not set(a_reduce_axes).intersection(a_var_axes), "%s: sources %r, red1 %r, red2 %r, var1 %r, var2 %r" % (
@@ -6404,6 +6407,25 @@ def __init__(self,
     res = tf.reshape(res, a_rem_dims + a_var_dims + b_var_dims)
     self.output.placeholder = res
 
+  @staticmethod
+  def _auto_var_axes(source1, source2, red1, red2):
+    """
+    :param Data source1:
+    :param Data source2:
+    :param list[int] red1:
+    :param list[int] red2:
+    :return: var1 tags, var2 tags
+    :rtype: (list[Dim], list[Dim])
+    """
+    is_equal_opts = dict(
+      treat_feature_as_spatial=True, allow_same_spatial_dim=True,
+      broadcast_matches=True, undefined_matches=True, derived_matches=True)
+    all_dim_tags, tags_dict = Dim.get_all_dimension_tags([source1, source2], is_equal_opts=is_equal_opts)
+    tags1, tags2 = tags_dict[source1], tags_dict[source2]
+    var1 = [tag for i, tag in enumerate(tags1) if tag not in tags2 and i not in red1]
+    var2 = [tag for i, tag in enumerate(tags2) if tag not in tags1 and i not in red2]
+    return var1, var2
+
   @staticmethod
   def _axis1_to_output(axis, a_rem_axes, a_var_axes):
     # Output will be of shape a_rem_dims + [a_var_dim, b_var_dim].
@@ -6440,29 +6462,30 @@ def transform_config_dict(cls, d, network, get_layer):
       # If the rec layer dim is the same as some other dim,
       # and was already explicitly specified in var1/var2 before,
       # skip it.
-      var1 = d.get("var1", -2)  # the default should really not be used...
-      var2 = d.get("var2", -1)  # the default should really not be used...
-      var1_ = set([src1.output.dim_tags[i] for i in src1.output.get_axes_from_description(var1)])
-      var2_ = set([src2.output.dim_tags[i] for i in src2.output.get_axes_from_description(var2)])
-      dims1.difference_update(var1_)
-      dims2.difference_update(var2_)
-      # The common dims should be shared. The shared common dims are implicit, so nothing to do about them.
-      dims_common = dims1.intersection(dims2)
-      # Those are dims which should be added to var1/var2.
-      dims1.difference_update(dims_common)
-      dims2.difference_update(dims_common)
-
-      def _add(dims, val, d_key):
-        if not dims:
-          return
-        if val is None or val == "":
-          val = []
-        elif not isinstance(val, (tuple, list)):
-          val = [val]
-        d[d_key] = val + type(val)(dims)
-
-      _add(dims1, var1, "var1")
-      _add(dims2, var2, "var2")
+      var1 = d.get("var1", NotSpecified)
+      var2 = d.get("var2", NotSpecified)
+      if var1 is not NotSpecified and var2 is not NotSpecified:
+        var1_ = set([src1.output.dim_tags[i] for i in src1.output.get_axes_from_description(var1)])
+        var2_ = set([src2.output.dim_tags[i] for i in src2.output.get_axes_from_description(var2)])
+        dims1.difference_update(var1_)
+        dims2.difference_update(var2_)
+        # The common dims should be shared. The shared common dims are implicit, so nothing to do about them.
+        dims_common = dims1.intersection(dims2)
+        # Those are dims which should be added to var1/var2.
+        dims1.difference_update(dims_common)
+        dims2.difference_update(dims_common)
+
+        def _add(dims, val, d_key):
+          if not dims:
+            return
+          if val is None or val == "":
+            val = []
+          elif not isinstance(val, (tuple, list)):
+            val = [val]
+          d[d_key] = val + type(val)(dims)
+
+        _add(dims1, var1, "var1")
+        _add(dims2, var2, "var2")
 
   @classmethod
   def get_out_data_from_opts(cls, name, sources,
@@ -6492,7 +6515,7 @@ def get_out_data_from_opts(cls, name, sources,
       message="DotLayer: Axes must be referenced by tag or special specified, not by int.",
       version=3)
     BehaviorVersion.require(
-      condition=all(a is not NotSpecified for a in (red1, red2, var1, var2)),
+      condition=all(a is not NotSpecified for a in (red1, red2)),
       message="DotLayer: Axes must be specified explicitly. There is no default.",
       version=3)
     BehaviorVersion.require(
@@ -6517,6 +6540,9 @@ def get_out_data_from_opts(cls, name, sources,
     assert not a_out.beam or not b_out.beam or a_out.beam == b_out.beam
     b_reduce_axes = b_out.get_axes_from_description(red2, allow_int=axis_desc_allow_int)
     assert a_reduce_axes and b_reduce_axes, "%s: sources %r, red1 %r, red2 %r" % (name, sources, red1, red2)
+    if BehaviorVersion.get() >= 3 and (var1 is NotSpecified or var2 is NotSpecified):
+      assert var1 is NotSpecified and var2 is NotSpecified
+      var1, var2 = cls._auto_var_axes(a_out, b_out, a_reduce_axes, b_reduce_axes)
     a_var_axes = a_out.get_axes_from_description(var1, allow_int=axis_desc_allow_int)
     b_var_axes = b_out.get_axes_from_description(var2, allow_int=axis_desc_allow_int)
     assert not set(a_reduce_axes).intersection(a_var_axes)


Test patch:
diff --git a/tests/test_TFNetworkRecLayer.py b/tests/test_TFNetworkRecLayer.py
--- a/tests/test_TFNetworkRecLayer.py
+++ b/tests/test_TFNetworkRecLayer.py
@@ -7359,8 +7359,7 @@ def _build_self_attention_layer(d, input, output, inside_rec_layer, query_axis,
   # Calculate the energies
   d[output + '_energy'] = {
     'class': 'dot', 'from': [output + '_query', output + '_key_accum'],
-    'red1': 'dim:%i' % key_dim, 'red2': 'dim:%i' % key_dim,
-    'var1': None if inside_rec_layer else query_axis, 'var2': key_dim_tag}  # [B,n,T?,T|rec-history]
+    'reduce': 'dim:%i' % key_dim}  # [B,n,T?,T|rec-history]
 
   d[output + '_weights'] = {
     'class': 'softmax_over_spatial', 'from': [output + '_energy'], 'axis': key_axis,
@@ -7371,8 +7370,7 @@ def _build_self_attention_layer(d, input, output, inside_rec_layer, query_axis,
 
   d[output + '_output'] = {
     'class': 'dot', 'from': [output + '_weights_drop', output + '_value_accum'],
-    'red1': key_axis, 'red2': key_axis,
-    "var1": None if inside_rec_layer else query_axis, "var2": "dim:%i" % value_dim}  # [B,n,T?,F|d_v]
+    'reduce': key_axis}  # [B,n,T?,F|d_v]
   d[output + '_att'] = {
     'class': 'merge_dims', 'axes': ["dim:%i" % num_heads, "dim:%i" % value_dim],
     'from': output + '_output'}  # [B,T?,F|n*d_v]