Repository: rwth-i6/returnn

Issue:
DotLayer: Support shuffled remaining dimensions
I have a config with a dot layer that takes two inputs of e.g. the shapes `[B,T,X,Y]` and `[T,B,Y,Z]`.
Then I want to compute their dot product, getting a tensor of the form `[B,T,X,Z]` (reduce the `Y` axis, have `X` and `Z` be the variable axes).

Note however that the the remaining axes (i.e. `B` and `T`) are equal but not in the same order for these two tensors.
Currently, this is not handled by `DotLayer` at all, it always assumes that all remaining axes are in the exact same order.
If they have static dimensions, than this will throw an error at template construction.
If not, the flag `debug=True` will assert that the shapes are compatible when the session is executed.
If that flag is not set, I guess TF would throw an error as well (except when it works because of broadcasting or sth like that).

Here is a small artificial example demonstrating this (with static shapes):
```
net_dict = {
  "a": {"class": "split_dims", "axis": "static:0", "dims": (2, 3, 5)},   # [B, 2, 3, 5]
  "b": {"class": "transpose", "from": ["a"], "perm": {"static:0": "static:1", "static:1": "static:0"}},  # [B, 3, 2, 5]
  "dot": {  # should produce [B, 2, 3, 1] or [B, 3, 2, 1], but throws an Exception instead
    "class": "dot", "from": ["a", "b"],
    "red1": "static:-1", "red2": "static:-1", "var1": None, "var2": None,
    "debug": True}
  # ...
}
```

To fix this, I think we need to transpose the remaining axes of one of the inputs to match each others order, which we can e.g.
determine by matching the `DimensionTag`s of the remaining axes using `DimensionTag.is_equal(..)` - similar to `get_common_data` does this as well.
But an extra `transpose` sounds inefficient (and from what I can tell from the code efficiency was considered thoroughly in the `DotLayer`) - perhaps we can integrate this better with the logic there already is in the layer?

Also: Why doesn't `DotLayer` use `tf.einsum` instead of `tf.matmul`? This would handle those index permutations by itself. (I guess performance? Or maybe backwards compatibility?)


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
@@ -4875,6 +4875,13 @@ def __init__(self, red1=-1, red2=-2, var1=-2, var2=-1, add_var2_if_empty=True, d
       self, self.sources, red1, red2, var1, var2)
     a_rem_axes = [i for i in range(a_out.batch_ndim) if i not in a_var_axes + a_reduce_axes]
     b_rem_axes = [i for i in range(b_out.batch_ndim) if i not in b_var_axes + b_reduce_axes]
+    assert len(a_rem_axes) == len(b_rem_axes)
+
+    # ensure that a_rem_axes and b_rem_axes are in the same order
+    map_a_to_b_rem_axes = b_out.find_matching_dim_map(a_out, a_rem_axes)
+    assert all(b_axis in map_a_to_b_rem_axes.values() for b_axis in b_rem_axes)
+    b_rem_axes = [map_a_to_b_rem_axes[a_axis] for a_axis in a_rem_axes]
+
     transpose_a = bool(a_var_axes and a_reduce_axes[0] < a_var_axes[0])
     transpose_b = bool(b_var_axes and b_reduce_axes[0] > b_var_axes[0])
     # For A, if not transpose_a, we must reorder the axes as: a_rem_axes + a_var_axes + a_reduce_axes.
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
@@ -112,7 +112,7 @@ def can_compare(self):
     return True
 
   def is_equal(self, other, ignore_feature_dim=False, allow_same_feature_dim=False, allow_same_spatial_dim=None,
-               treat_feature_as_spatial=False, broadcast_matches=False):
+               treat_feature_as_spatial=False, broadcast_matches=False, unknown_spatial_matches=False):
     """
     Compares self to other for equality.
     Note that the default behavior is very restrictive.
@@ -125,6 +125,7 @@ def is_equal(self, other, ignore_feature_dim=False, allow_same_feature_dim=False
     :param bool|None allow_same_spatial_dim:
     :param bool treat_feature_as_spatial:
     :param bool broadcast_matches:
+    :param bool unknown_spatial_matches:
     :rtype: bool
     """
     if allow_same_spatial_dim is None:
@@ -161,6 +162,8 @@ def is_equal(self, other, ignore_feature_dim=False, allow_same_feature_dim=False
     if self_kind == other_kind == self.Types.Spatial:
       if self.dimension is not None and allow_same_spatial_dim:
         return True
+      if unknown_spatial_matches and ((self.dyn_size is None) != (other.dyn_size is None)):
+        return True
     if self.description == other.description:
       return True
     return False
@@ -2501,3 +2504,56 @@ def get_common_data(cls, sources, warnings_out=None, out_shape=None):
     # Simple fallback: Use first with biggest batch_ndim.
     # Was even simpler before: Use first.
     return common
+
+  def find_matching_dims(self, dim_tag, is_equal_opts):
+    """
+    Finds the dimensions of this Data that match another DimensionTag
+
+    :param DimensionTag dim_tag:
+    :param dict[str,bool]|None is_equal_opts: passed to DimensionTag.is_equal
+    :rtype: list[int] a list of matching axes, counted with batch dim. Sorted in ascending order
+    """
+    return [axis for axis in range(self.batch_ndim) if self.get_dim_tag(axis).is_equal(dim_tag, **is_equal_opts)]
+
+  def find_matching_dim_map(self, other, other_axes):
+    """
+    Looks up all other_axes of another Data in this Data. Does not allow duplicates.
+
+    :param Data other:
+    :param list[int] other_axes: a list of axes of ``other``, counted with batch dim
+    :return: a dict mapping other axes to own axes, all counted with batch dim
+    :rtype: dict[int,int]
+    """
+    def map_other_axis_to_self(other_axis, taken_self_axes):
+      """
+      :param int other_axis: counted with batch dim
+      :param list[int] taken_self_axes: axes that should not be used again
+      :return: the axis of ``self`` that matches ``other_axis``, counted with batch dim
+      :rtype: int
+      """
+      is_equal_opts = dict(
+        allow_same_feature_dim=True, allow_same_spatial_dim=True, treat_feature_as_spatial=True)
+      other_axis_dim_tag = other.get_dim_tag(other_axis)
+      matching = [
+        self_axis for self_axis in self.find_matching_dims(other_axis_dim_tag, is_equal_opts)
+        if self_axis not in taken_self_axes
+      ]
+      if not matching:
+        # The DimensionTags do not match. Then we also allow one single dyn_size to be unknown
+        is_equal_opts["unknown_spatial_matches"] = True
+        matching = [
+          self_axis for self_axis in self.find_matching_dims(other_axis_dim_tag, is_equal_opts)
+          if self_axis not in taken_self_axes
+        ]
+        assert len(matching) == 1, 'cannot match the axes %s from %s to %s' % (other_axes, other, self)
+      assert matching, 'cannot match the axes %s from %s to %s' % (other_axes, other, self)
+      # 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]
+
+    other_to_self_mapping = {}
+    for axis in other_axes:
+      other_to_self_mapping[axis] = map_other_axis_to_self(axis, other_to_self_mapping.values())
+    assert len(other_to_self_mapping) == len(other_axes), 'other_axes may not contain duplicates'
+
+    return other_to_self_mapping


Test patch:
diff --git a/tests/test_TFNetworkLayer.py b/tests/test_TFNetworkLayer.py
--- a/tests/test_TFNetworkLayer.py
+++ b/tests/test_TFNetworkLayer.py
@@ -723,6 +723,64 @@ def test_CombineLayer_two_time_dims_first_not_most_generic_with_n_out():
     assert out_np.shape == (n_time1, n_batch, n_time2, n_dim)
 
 
+def test_dot_layer_shuffled_remaining_dims_static():
+  with make_scope() as session:
+    import numpy as np
+    net_dict = {
+      "a": {"class": "split_dims", "axis": "static:0", "dims": (2, 3, 5)},
+      "b": {"class": "transpose", "from": ["a"], "perm": {"static:0": "static:1", "static:1": "static:0"}},
+      "dot": {
+        "class": "dot", "from": ["a", "b"],
+        "red1": "static:-1", "red2": "static:-1", "var1": None, "var2": None,
+        "debug": True},
+      "output": {"class": "merge_dims", "axes": "static", "from": ["dot"]}
+    }
+    config = Config()
+    config.update({
+      "extern_data": {"data": {"shape": (30,)}},
+      "network": net_dict
+    })
+    network = TFNetwork(config=config, train_flag=True)
+    network.construct_from_dict(config.typed_dict["network"])
+    out = network.get_default_output_layer(must_exist=True)
+    input_data = np.ones(shape=(17, 30))
+    feed_dict = {network.layers['data'].output.placeholder: input_data}
+
+    # just check that it runs
+    session.run(out.output.placeholder, feed_dict)
+
+
+def test_dot_layer_shuffled_remaining_dims_dynamic():
+  with make_scope() as session:
+    import numpy as np
+    batch_size = 8
+    time_size = 20
+    feat_size = 10
+    net_dict = {
+      "a": {"class": "copy", "from": "data"},
+      "b": {"class": "transpose", "from": ["a"], "perm": {"B": "T", "T": "B"}},
+      "dot": {
+        "class": "dot", "from": ["a", "b"],
+        "red1": "F", "red2": "F", "var1": None, "var2": None,
+        "debug": True},
+      "output": {"class": "copy", "from": ["dot"]}
+    }
+    config = Config()
+    config.update({
+      "num_outputs": 1,
+      "num_inputs": feat_size,
+      "network": net_dict
+    })
+    network = TFNetwork(config=config, train_flag=True)
+    network.construct_from_dict(config.typed_dict["network"])
+    out = network.get_default_output_layer(must_exist=True)
+    input_data = np.ones(shape=(batch_size, time_size, feat_size))
+    feed_dict = {network.layers['data'].output.placeholder: input_data}
+
+    # just check that it runs
+    session.run(out.output.placeholder, feed_dict)
+
+
 def test_dropout_layer_net_construct():
   with make_scope() as session:
     net_dict = {