Repository: rwth-i6/returnn

Issue:
DotLayer masking heuristic fails in some cases with `dyn_size_ext`
DotLayer currently has this heuristic to determine on which input we apply masking:

```
        # We need to apply a mask.
        # We don't need it on both a and b. We can either apply it on a or on b.
        # Use some very simple heuristic where the mask is maybe cheaper.
        if len(a_shape) < len(b_shape):
          a = mask_dyn_seq_len_nd(a_out, pad_value=0, axes=a_reduce_dyn_axes)
          self._info_reduce_mask = "mask-source-0"
        else:
          b = mask_dyn_seq_len_nd(b_out, pad_value=0, axes=b_reduce_dyn_axes)
          self._info_reduce_mask = "mask-source-1"
```

However with `dyn_size_ext`, it can happen that the seq lengths depend on another axis which is only present in one of these inputs.
In our case:
```
a: Data{'att_weights', [B?,F|'heads'(8),T|'classes'[B?],'keys'['classes'[B?]]]}
b: Data{'att_value', [T|'keys'['classes'[B?]],B?,'heads'(8),F|F'feature'(64)]}
```
the seq lengths of `keys` depends on `classes`, so we have to mask input `a`.
DotLayer masking heuristic fails in some cases with `dyn_size_ext`
DotLayer currently has this heuristic to determine on which input we apply masking:

```
        # We need to apply a mask.
        # We don't need it on both a and b. We can either apply it on a or on b.
        # Use some very simple heuristic where the mask is maybe cheaper.
        if len(a_shape) < len(b_shape):
          a = mask_dyn_seq_len_nd(a_out, pad_value=0, axes=a_reduce_dyn_axes)
          self._info_reduce_mask = "mask-source-0"
        else:
          b = mask_dyn_seq_len_nd(b_out, pad_value=0, axes=b_reduce_dyn_axes)
          self._info_reduce_mask = "mask-source-1"
```

However with `dyn_size_ext`, it can happen that the seq lengths depend on another axis which is only present in one of these inputs.
In our case:
```
a: Data{'att_weights', [B?,F|'heads'(8),T|'classes'[B?],'keys'['classes'[B?]]]}
b: Data{'att_value', [T|'keys'['classes'[B?]],B?,'heads'(8),F|F'feature'(64)]}
```
the seq lengths of `keys` depends on `classes`, so we have to mask input `a`.


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
@@ -5626,12 +5626,18 @@ def __init__(self, red1=-1, red2=-2, var1=-2, var2=-1, add_var2_if_empty=True, d
         self._info_reduce_mask = "source-1-already-masked"  # it's already masked as needed
       else:
         # We need to apply a mask.
-        # We don't need it on both a and b. We can either apply it on a or on b.
-        # Use some very simple heuristic where the mask is maybe cheaper.
-        if len(a_shape) < len(b_shape):
+        # We don't need it on both a and b. In case we can either apply it on a or on b,
+        # use some very simple heuristic where the mask is maybe cheaper.
+        can_mask_a = all(
+          set(a_out.dim_tags[i].dyn_size_ext.dim_tags).issubset(a_out.dim_tags) for i in a_reduce_dyn_axes)
+        can_mask_b = all(
+          set(b_out.dim_tags[i].dyn_size_ext.dim_tags).issubset(b_out.dim_tags) for i in b_reduce_dyn_axes)
+        if not can_mask_b or len(a_shape) < len(b_shape):
+          assert can_mask_a
           a = mask_dyn_seq_len_nd(a_out, pad_value=0, axes=a_reduce_dyn_axes)
           self._info_reduce_mask = "mask-source-0"
         else:
+          assert can_mask_b
           b = mask_dyn_seq_len_nd(b_out, pad_value=0, axes=b_reduce_dyn_axes)
           self._info_reduce_mask = "mask-source-1"
     else:


Test patch:
diff --git a/tests/test_TFNetworkLayer.py b/tests/test_TFNetworkLayer.py
--- a/tests/test_TFNetworkLayer.py
+++ b/tests/test_TFNetworkLayer.py
@@ -3738,6 +3738,31 @@ def test_DotLayer_mask_dyn_seq_after_softmax():
     session.run(layer.output.placeholder, feed_dict=feed_dict)
 
 
+def test_DotLayer_self_att_dyn_size_ext():
+  batch_dim = DimensionTag(kind=DimensionTag.Types.Batch)
+  heads_dim = DimensionTag(kind=DimensionTag.Types.Spatial, description="heads", dimension=8)
+  classes_dim = DimensionTag(kind=DimensionTag.Types.Time, description="classes")
+  keys_dim = DimensionTag(
+    kind=DimensionTag.Types.Spatial, description="keys",
+    dyn_size_ext=Data(name="keys_dyn_size", dim_tags=[classes_dim], dtype="int32", auto_create_placeholders=True))
+  feature_dim = DimensionTag(kind=DimensionTag.Types.Feature, description="feature", dimension=64)
+
+  a = Data(name="att_weights", dim_tags=[batch_dim, heads_dim, classes_dim, keys_dim], auto_create_placeholders=True)
+  b = Data(name="att_value", dim_tags=[keys_dim, batch_dim, heads_dim, feature_dim], auto_create_placeholders=True)
+  print("a:", a)
+  print("b:", b)
+
+  config = Config({"debug_print_layer_output_template": True})
+  config.update(dict(num_inputs=64 * 8, num_outputs=64 * 8))
+  net = TFNetwork(config=config, train_flag=True, search_flag=False, eval_flag=False)
+
+  a_lay = InternalLayer(name="att_weights", network=net, output=a)
+  b_lay = InternalLayer(name="att_value", network=net, output=b)
+  dot_kwargs = {"red1": keys_dim, "red2": keys_dim, "var1": classes_dim, "var2": feature_dim}
+  dot = DotLayer.get_out_data_from_opts(name="dot", sources=[a_lay, b_lay], **dot_kwargs)
+  DotLayer(network=net, name="dot", sources=[a_lay, b_lay], output=dot, **dot_kwargs)  # just check that it builds.
+
+
 def test_subnet_load_on_init():
   import tempfile
   model_tmp_dir = tempfile.mkdtemp("tmp-checkpoint")