Repository: rwth-i6/returnn

Issue:
`MaskedComputationLayer` is violating the principle that the user should not need to think about rec automatic optimization
See [RETURNN principles](https://github.com/rwth-i6/returnn/wiki/RETURNN-principles).

When defining `MaskedComputationLayer` inside a rec loop and not thinking about rec automatic optimization, the behavior is clear. It copies the prev output (and prev state) whenever the mask is False, otherwise it uses the current output.

In this line of thought, there should be no need for `UnmaskLayer`.

However, the current situation is, when it is optimized out, a tensor of shape [B,D] inside would not just become [T,B,D] where T is the number of rec iterations, but instead it becomes [T',B,D] where T' = sum(mask). And some follow-up operation at some point requires the [T,B,D] shape (maybe because it combines it with other tensors from the rec loop), and thus the user explicitly must use the `UnmaskLayer` at some point after the `MaskedComputationLayer`.

This violates the RETURNN principle that the user should not need to think about rec automatic optimization.

---

What are possible solutions?

I don't really have a good idea so far. It should be opaque.

Maybe the new dim tag for T' could have some special flag that it is a masked subset of T, and whenever it is going to be combined with the other dim tag, it would automatically do the unmasking.



Discussion:
@Zettelkasten @mmz33 @robin-p-schmitt maybe you have some ideas, or comments?
Note that I already realized this as a problem some time ago, but now again when thinking about specifying the output spatial dim (`out_spatial_dim` #597) and the exact behavior.

When not being optimized (e.g. simply `optimize_move_layers_out=False`) and then accumulating the output of `MaskedComputationLayer`, you would get shape [T,B,D], and never actually the shape [T',B,D].

`out_spatial_dim` should in all cases refer to T', even when T' would actually never exist? So basically `out_spatial_dim` would be ignored inside the loop, and only used outside a rec loop?

The new dim T' would somehow have a reference to T. `derived_from_tag` is obvious, but we need more, esp we need the mask, or the indices for unmasking.

It's a bit like when e.g. `ConvLayer` did downsampling via striding. Then T' = T / 2. And this could be formulated as a mask as well.

Now we can see the unmasking as a sparse-to-dense operation.

However, this is still ambiguous. In sparse-to-dense, we would set 0 for all other frames. But this is rarely what we would want here. Specifically, when reproducing the output of a `MaskedComputationLayer`, we would have a very specific behavior, namely that the previous masked frame would be copied.

So, when we want to have automatic unmasking of T' to T, when such tensors are combined somewhere, e.g. via `Data.get_common_data` and then `Data.copy_compatible_to`, we don't just need the mask but we also need the kind of unmasking, i.e. how to fill the other frames.

We could use the `derived_from_op` mechanism, via a special op kind `"mask"`, and store the mask in the attributes, along with `unmask_type="left"` or so, to reflect that we want this specific behavior when unmasking.

This could later be extended to be able to represent sparse data via `unmask_type="fill"` and `unmask_fill_value`.

The mask itself would be a `Data` instance. In our case, it would be of shape [B,T] with dtype bool. However, the shape is probably arbitrary. Although it probably should contain the original dim T.

Or as mentioned, instead of a mask, we could also store the indices, i.e. `Data` of shape [B,T'] pointing to T.


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
@@ -3126,7 +3126,16 @@ def __init__(self, axis="T", additional_left_summand_per_element=None, reverse=F
       self.output.placeholder = next_state
     else:
       axis = data.get_axis_from_description(axis)
-      self.output.placeholder = tf.cumsum(x, axis=axis, reverse=reverse)
+      y = tf.cumsum(x, axis=axis, reverse=reverse)
+      if self._initial_output is not None:
+        init_state = self.get_rec_initial_output(
+          name=self.name, output=self.output,
+          network=self.network, batch_dim=self.get_batch_dim(),
+          rec_layer=self.network.get_rec_parent_layer(inside_loop=False),  # should actually not matter
+          axis=axis, sources=self.sources,
+          initial_output=self._initial_output)
+        y = init_state + y
+      self.output.placeholder = y
     self.output.placeholder.set_shape(data.placeholder.get_shape())
     self.output.placeholder.set_shape(tf.TensorShape(self.output.batch_shape))
     self.output.size_placeholder = self.input_data.size_placeholder.copy()
@@ -3157,7 +3166,9 @@ def get_rec_initial_extra_outputs(cls, network, batch_dim, rec_layer, axis="T",
     from returnn.tf.util.basic import is_axis_from_description_recurrent
     if is_axis_from_description_recurrent(axis=axis, network=network, data=data):
       assert all(data.shape)
-      return {"state": tf.zeros(data.get_batch_shape(batch_dim=batch_dim), dtype=data.dtype)}
+      init_state = cls.get_rec_initial_output(
+        network=network, batch_dim=batch_dim, rec_layer=rec_layer, axis=axis, sources=sources, **kwargs)
+      return {"state": init_state}
     return {}
 
 
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
@@ -7578,6 +7578,7 @@ class MaskedComputationLayer(LayerBase):
 
   def __init__(self, mask, unit, masked_from,
                _layer_class, _layer_desc,
+               in_spatial_dim=None,
                out_spatial_dim=None,
                _parent_layer_cache=None,
                **kwargs):
@@ -7585,12 +7586,12 @@ def __init__(self, mask, unit, masked_from,
     :param LayerBase|None mask:
     :param dict[str] unit:
     :param LayerBase|None masked_from:
-    :param Dim|None out_spatial_dim:
+    :param Dim|None in_spatial_dim:
+    :param Dim|None out_spatial_dim: the masked dim
     :param type[LayerBase] _layer_class:
     :param dict[str] _layer_desc:
     :param dict[str,LayerBase]|None _parent_layer_cache:
     """
-    out_spatial_dim  # noqa  # handled in transform_config_dict
     from returnn.tf.network import get_layer_class
     from .base import WrappedInternalLayer
     from returnn.tf.util.basic import where_bc, get_shape, nd_indices
@@ -7611,16 +7612,15 @@ def __init__(self, mask, unit, masked_from,
       else:
         assert mask.output.have_time_axis() and mask.output.shape == (None,) and mask.output.dtype == "bool", (
           "%s: invalid mask %s (outside rec loop)" % (self, mask))
+        assert in_spatial_dim and out_spatial_dim
         mask_data = mask.output.copy_as_time_major()
         mask_t = where_bc(mask_data.placeholder, mask_data.get_sequence_mask(), tf.convert_to_tensor(False))
         idxs = tf.cumsum(tf.cast(mask_t, tf.int32), axis=0)  # [T,B] -> idx in T' + 1
         if masked_from:
           new_size = masked_from.output.get_sequence_lengths()
-          new_dim_tag = masked_from.output.get_time_dim_tag()
         else:
           new_size = idxs[-1]  # [B]
-          new_dim_tag = self.output.get_time_dim_tag()
-          new_dim_tag.dyn_size = new_size
+          out_spatial_dim.dyn_size = new_size
         new_time = tf.reduce_max(new_size)  # T'
         idxs = where_bc(mask_t, idxs - 1, new_time)
 
@@ -7636,12 +7636,17 @@ def get_masked_layer(source):
         # We can just leave it as-is. The state will handled below.
         return source
       else:
-        source_data = source.output.copy_as_time_major()
+        assert in_spatial_dim
+        if in_spatial_dim not in source.output.dim_tags:
+          return source
+        axis = source.output.get_axis_from_description(in_spatial_dim)
+        source_data = source.output.copy_move_axis(old_axis=axis, new_axis=0)  # time-major
+        source_data.time_dim_axis = 0
         assert source_data.is_same_time_dim(mask_data), "%s mask and source time dim do not match" % self
         tmp_shape = get_shape(source_data.placeholder)
         tmp_shape[0] = new_time + 1  # one more for the padded data
         res = tf.scatter_nd(nd_indices(idxs, batch_axis=1), source_data.placeholder, shape=tmp_shape)
-        res_data = source_data.copy_template().copy_template_replace_dim_tag(axis=0, new_dim_tag=new_dim_tag)
+        res_data = source_data.copy_template().copy_template_replace_dim_tag(axis=0, new_dim_tag=out_spatial_dim)
         res_data.placeholder = res[:new_time]
         res_data.beam = SearchBeam.get_combined_beam(res_data.beam, mask.output.beam)
         layer_desc = dict(base_layer=source, network=source.network, name=source.name, output=res_data)
@@ -7661,7 +7666,7 @@ def get_masked_layer(source):
       source.sources.append(masked_from)  # add dep
       sub_layers["data"] = source
 
-    else:
+    elif len(self.sources) >= 1:
       assert len(self.sources) == 1
       sub_layers["data"] = get_masked_layer(self.sources[0])
 
@@ -7692,13 +7697,30 @@ def sub_get_layer(sub_layer_name):
     layer_class.transform_config_dict(layer_desc, network=extra_net, get_layer=sub_get_layer)
     # noinspection PyProtectedMember
     layer_desc = extra_net._create_layer_layer_desc(name=self.name, layer_desc=layer_desc)
-    layer_desc["output"] = self.output.copy_template(name="%s_output" % self.name)
+    sub_out_template = self.output.copy_template(name="%s_output" % self.name)
+    if in_spatial_dim and in_spatial_dim in sub_out_template.dim_tags:  # could happen when optimized
+      assert out_spatial_dim
+      axis = sub_out_template.get_axis_from_description(in_spatial_dim)
+      sub_out_template = sub_out_template.copy_template_replace_dim_tag(axis=axis, new_dim_tag=out_spatial_dim)
+    layer_desc["output"] = sub_out_template
 
     self.sub_layer = layer_class(**layer_desc)
     assert isinstance(self.sub_layer, LayerBase)
     self.sub_layer.post_init(layer_desc)
     self.sub_layer.output.sanity_check()
-    self.output = self.sub_layer.output.copy(name="%s_output" % self.name)
+    inside_rec_time_dim = self.network.get_inside_rec_time_dim(inside_loop=True)
+    over_rec_time_dim = self.network.get_inside_rec_time_dim(inside_loop=False)
+    self._unmask_layer = None
+    # In case we are in rec layer and optimize out of loop and mask over this axis, and don't have masked_from:
+    if not masked_from and over_rec_time_dim == in_spatial_dim and over_rec_time_dim and not inside_rec_time_dim:
+      # We will automatically unmask again (https://github.com/rwth-i6/returnn/pull/976).
+      assert self.mask
+      self._unmask_layer = UnmaskLayer(
+        name="%s(internal-unmask)" % self.name, mask=self.mask, sources=[self.sub_layer], network=self.network,
+        output=self.output.copy_template())
+      self.output = self._unmask_layer.output.copy(name="%s_output" % self.name)
+    else:
+      self.output = self.sub_layer.output.copy(name="%s_output" % self.name)
     self.rec_vars_outputs = self.sub_layer.rec_vars_outputs.copy()
     self.params = self.sub_layer.params
 
@@ -7773,10 +7795,11 @@ def transform_config_dict(cls, d, network, get_layer):
     super(MaskedComputationLayer, cls).transform_config_dict(d, network=network, get_layer=get_layer)
     # Just call it for dep resolution.
     parent_layer_cache = d.setdefault("_parent_layer_cache", {})
-    d["_layer_class"], d["_layer_desc"] = cls._create_template(
+    d["_layer_class"], d["_layer_desc"], d["in_spatial_dim"], d["out_spatial_dim"] = cls._create_template(
       name=d["_name"], network=network, sources=d["sources"],
       masked_from=masked_from,
       unit=d["unit"],
+      in_spatial_dim=d.get("in_spatial_dim", None),
       out_spatial_dim=d.get("out_spatial_dim", None),
       get_layer=get_layer, _parent_layer_cache=parent_layer_cache)
     if masked_from and not parent_layer_cache:
@@ -7788,7 +7811,7 @@ def transform_config_dict(cls, d, network, get_layer):
   # noinspection PyUnusedLocal
   @classmethod
   def _create_template(cls, name, network, sources, masked_from, unit,
-                       out_spatial_dim=None,
+                       in_spatial_dim=None, out_spatial_dim=None,
                        get_layer=None, _parent_layer_cache=None, **kwargs):
     """
     :param str name:
@@ -7796,74 +7819,91 @@ def _create_template(cls, name, network, sources, masked_from, unit,
     :param list[LayerBase] sources:
     :param LayerBase masked_from:
     :param dict[str] unit:
-    :param Dim|None out_spatial_dim:
+    :param Dim|None in_spatial_dim:
+    :param Dim|None out_spatial_dim: the masked dim
     :param (str)->LayerBase get_layer:
     :param dict[str,LayerBase]|None parent_layer_cache:
-    :return: layer_class, layer_desc
+    :return: layer_class, layer_desc, in_spatial_dim, out_spatial_dim
     """
     from returnn.tf.network import get_layer_class
     from .base import WrappedInternalLayer
     if not get_layer:
       get_layer = network.get_layer
-    # We don't care about the right masked input here, but just about deriving the right output shape.
-    if masked_from:
-      if out_spatial_dim:
-        masked_from.output.get_time_dim_tag().declare_same_as(out_spatial_dim)
-      if network.is_inside_rec_layer(inside_loop=True):
-        source_data = (
-          masked_from.output
-          .copy_template_excluding_time_dim(
-            name="%s_%s_masked_input_frame" % (masked_from.output.name, name))
-          .copy_template_set_ctx(network.get_control_flow_ctx()))
-      else:
-        source_data = masked_from.output.copy_template(
-          name="%s_%s_masked_input" % (masked_from.output.name, name))
-      source_data.available_for_inference = True  # we would make sure that this works at inference
-      source = WrappedInternalLayer(
-        base_layer=masked_from, network=masked_from.network, name=masked_from.name, output=source_data)
-    else:
-      assert len(sources) == 1
-      source, = sources
-      assert isinstance(source, LayerBase) or not source
-      if not network.is_inside_rec_layer() and source:
-        source_data = source.output.copy_template().copy_as_time_major()
-        # Create own time dim tag, to make sure we have some own custom.
-        if not out_spatial_dim:
-          out_spatial_dim = Dim(
-            kind=Dim.Types.Spatial, description="%s:masked:time" % name,
-            derived_from_tag=source_data.get_time_dim_tag(), auto_generated=True)
-        source_data = source_data.copy_template_replace_dim_tag(
-          axis=0,
-          new_dim_tag=out_spatial_dim)
-        source = WrappedInternalLayer(
-          base_layer=source, network=source.network, name=source.name,
-          output=source_data)
+
+    over_rec_time_dim = network.get_inside_rec_time_dim(inside_loop=False)
+    if over_rec_time_dim and not in_spatial_dim:
+      in_spatial_dim = over_rec_time_dim
+
+    class _Locals:  # Python 2, and modify this maybe via sub_get_layer
+      in_spatial_dim_ = in_spatial_dim
+      out_spatial_dim_ = out_spatial_dim
+
+    _layer_cache = {}
 
     def sub_get_layer(sub_layer_name):
       """
       :param str sub_layer_name:
       :rtype: LayerBase
       """
-      if sub_layer_name == "data":
+      if sub_layer_name in _layer_cache:
+        return _layer_cache[sub_layer_name]
+
+      if sub_layer_name == "data" and masked_from:
+        if _Locals.out_spatial_dim_:
+          if _Locals.out_spatial_dim_ not in masked_from.output.dim_tags:
+            masked_from.output.get_time_dim_tag().declare_same_as(_Locals.out_spatial_dim_)
+        else:
+          _Locals.out_spatial_dim_ = masked_from.output.get_time_dim_tag()
+        if network.is_inside_rec_layer():
+          source_data = (
+            masked_from.output
+            .copy_template_excluding_time_dim(
+              name="%s_%s_masked_input_frame" % (masked_from.output.name, name))
+            .copy_template_set_ctx(network.get_control_flow_ctx()))
+        else:
+          source_data = masked_from.output.copy_template(
+            name="%s_%s_masked_input" % (masked_from.output.name, name))
+        source_data.available_for_inference = True  # we would make sure that this works at inference
+        source = WrappedInternalLayer(
+          base_layer=masked_from, network=masked_from.network, name=masked_from.name, output=source_data)
+        _layer_cache[sub_layer_name] = source
         return source
-      if _parent_layer_cache and sub_layer_name in _parent_layer_cache:
+
+      elif sub_layer_name == "data" and len(sources) >= 1:
+        assert len(sources) == 1
+        layer = sources[0]
+      elif _parent_layer_cache and sub_layer_name in _parent_layer_cache:
         layer = _parent_layer_cache[sub_layer_name]
       else:
         assert not sub_layer_name.startswith(extra_net.extra_name_prefix + ":")
         layer = get_layer(sub_layer_name)
         if not layer:
-          return layer
+          return None
         if _parent_layer_cache is not None:
           _parent_layer_cache[sub_layer_name] = layer
-      if not network.is_inside_rec_layer():
-        # noinspection PyShadowingNames
-        source_data_ = layer.output.copy_template().copy_as_time_major()
-        source_data_ = source_data_.copy_template_replace_dim_tag(axis=0, new_dim_tag=source.output.get_time_dim_tag())
+      source_data = layer.output.copy_template()
+      if not _Locals.in_spatial_dim_:
+        _Locals.in_spatial_dim_ = source_data.get_time_dim_tag()
+
+      if _Locals.in_spatial_dim_ in source_data.dim_tags:
+        axis = source_data.get_axis_from_description(_Locals.in_spatial_dim_)
+        source_data = source_data.copy_move_axis(old_axis=axis, new_axis=0)  # time-major
+        # Create own time dim tag, to make sure we have some own custom.
+        if not _Locals.out_spatial_dim_:
+          _Locals.out_spatial_dim_ = Dim(
+            kind=Dim.Types.Spatial, description="%s:masked:time" % name,
+            derived_from_tag=_Locals.in_spatial_dim_, auto_generated=True)
+        source_data = source_data.copy_template_replace_dim_tag(axis=0, new_dim_tag=_Locals.out_spatial_dim_)
         layer = WrappedInternalLayer(
           base_layer=layer, network=layer.network, name=layer.name,
-          output=source_data_)
+          output=source_data)
+      _layer_cache[sub_layer_name] = layer
       return layer
 
+    if sources or masked_from:
+      # Make sure to resolve this first, to maybe declare in_spatial_dim/out_spatial_dim.
+      sub_get_layer("data")
+
     extra_net = network.make_extra_net(
       prefix_name="extra._internal_template.masked(%s)" % name,
       # There should be no need by the base get_layer to get back to us.
@@ -7879,22 +7919,34 @@ def sub_get_layer(sub_layer_name):
     layer_class.transform_config_dict(layer_desc, network=extra_net, get_layer=sub_get_layer)
     # noinspection PyProtectedMember
     layer_desc = extra_net._create_layer_layer_desc(name=name, layer_desc=layer_desc)
-    return layer_class, layer_desc
+    return layer_class, layer_desc, _Locals.in_spatial_dim_, _Locals.out_spatial_dim_
 
   @classmethod
-  def get_out_data_from_opts(cls, network, out_spatial_dim=None, **kwargs):
+  def get_out_data_from_opts(cls, network, masked_from=None, in_spatial_dim=None, out_spatial_dim=None, **kwargs):
     """
     :param returnn.tf.network.TFNetwork network:
+    :param LayerBase|Noen masked_from:
+    :param Dim|None in_spatial_dim:
     :param Dim|None out_spatial_dim:
     :rtype: Data
     """
-    out_spatial_dim  # noqa  # handled in transform_config_dict
     layer_class, layer_desc = kwargs["_layer_class"], kwargs["_layer_desc"]
+    assert issubclass(layer_class, LayerBase)
     output = layer_class.get_out_data_from_opts(**layer_desc)
     assert isinstance(output, Data)
     output.sanity_check()
+    if out_spatial_dim and out_spatial_dim in output.dim_tags:
+      axis = output.get_axis_from_description(out_spatial_dim)
+      output = output.copy_template_replace_dim_tag(axis=axis, new_dim_tag=out_spatial_dim)
+    inside_rec_time_dim = network.get_inside_rec_time_dim(inside_loop=True)
+    over_rec_time_dim = network.get_inside_rec_time_dim(inside_loop=False)
     if network.is_inside_rec_layer():
       output = output.copy_as_batch_major()
+    elif not masked_from and in_spatial_dim == over_rec_time_dim and over_rec_time_dim and not inside_rec_time_dim:
+      # We will automatically unmask again (https://github.com/rwth-i6/returnn/pull/976).
+      assert out_spatial_dim and out_spatial_dim in output.dim_tags
+      axis = output.get_axis_from_description(out_spatial_dim)
+      output = output.copy_template_replace_dim_tag(axis=axis, new_dim_tag=in_spatial_dim)
     return output
 
   def get_constraints_value(self):
@@ -7959,9 +8011,11 @@ def get_rec_initial_extra_outputs(cls, batch_dim, rec_layer, **kwargs):
     assert isinstance(output, Data)
     assert issubclass(layer_class, LayerBase)
     with layer_class.cls_setup_scope(**kwargs):
-      d = layer_class.get_rec_initial_extra_outputs(batch_dim=batch_dim, rec_layer=rec_layer, **layer_desc)
+      output = output.copy_as_batch_major()
+      d = layer_class.get_rec_initial_extra_outputs(
+        batch_dim=batch_dim, rec_layer=rec_layer, output=output, **layer_desc)
       initial_out = layer_class.get_rec_initial_output(
-        batch_dim=batch_dim, rec_layer=rec_layer, output=output.copy_as_batch_major(), **layer_desc)
+        batch_dim=batch_dim, rec_layer=rec_layer, output=output, **layer_desc)
       assert "_output" not in d
       d["_output"] = initial_out
       return d
@@ -8012,6 +8066,7 @@ def __init__(self, mask, **kwargs):
     self.mask = mask
     src_layer = self.sources[0]
     batch_dim = self.get_batch_dim()
+    orig_time_dim = self.mask.output.get_time_dim_tag() if self.mask.output.have_time_axis() else None
     if not src_layer.output.have_time_axis():
       assert self.network.is_inside_rec_layer()
       assert self.output.placeholder is not None  # should be the copy of source already
@@ -8019,12 +8074,15 @@ def __init__(self, mask, **kwargs):
       if self.network.is_inside_rec_layer():
         # We have this state, although we don't need it, we still must set it.
         self.rec_vars_outputs["t"] = tf.zeros([batch_dim], dtype=tf.int32) - 1
-    else:
+
+    else:  # src has time dim
       rec_parent_layer = self.network.get_rec_parent_layer(inside_loop=False)
       assert rec_parent_layer and isinstance(rec_parent_layer.cell, _SubnetworkRecCell)
       # noinspection PyProtectedMember
       initial = rec_parent_layer.cell._get_init_output(src_layer.name, batch_dim=batch_dim)  # [B,D']
       if self.network.is_inside_rec_layer():
+        # This UnmaskLayer is inside the rec loop but the source is outside (or at least does not have a time dim).
+        # The RecLayer will not unroll the source when the dim tag do not match, i.e. when it is masked.
         with same_control_flow_ctx(src_layer.output.placeholder):
           src = src_layer.output.copy_as_bt_or_tb_major()
         mask_out = self.mask.output
@@ -8038,6 +8096,11 @@ def __init__(self, mask, **kwargs):
         y = where_bc(tf.equal(t, -1)[:, None], initial, y)
         self.output.placeholder = y
 
+      elif orig_time_dim in src_layer.output.dim_tags:  # outside rec loop, and src seems already unmasked
+        # Do nothing, as src already seems unmasked.
+        # Note that we might run into this case due to: https://github.com/rwth-i6/returnn/pull/976
+        self.output.placeholder = src_layer.output.get_placeholder_as_time_major()
+
       else:  # outside rec loop
         assert src_layer.output.get_time_dim_tag() != self.mask.output.get_time_dim_tag(), (
           "%s: unexpected source" % self)
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
@@ -841,6 +841,24 @@ def get_same_derived_base(self):
       assert base
     return base
 
+  def get_derived_bases_list(self):
+    """
+    :rtype: Dim
+    """
+    res = [self]
+    base = self
+    visited = {}
+    while base.same_as or base.derived_from_tag:
+      assert id(base) not in visited  # should not have cycles. normally this should never be triggered
+      visited[id(base)] = base
+      if base.same_as:
+        base = base.same_as
+        continue
+      base = base.derived_from_tag
+      assert base
+      res.append(base)
+    return res
+
   @property
   def undefined(self):
     """
@@ -997,8 +1015,10 @@ def get_all_dimension_tags(cls, data_list, is_equal_opts=None, unique_separate_a
         if existing_tag:
           if unique_separate_axes:
             existing_tag_collection_for_data.remove(existing_tag)  # don't take it again for this data
-          if existing_tag.undefined and not tag.undefined and tag.dimension == existing_tag.dimension:
-            # Replace the existing by the new tag.
+          replace_existing = existing_tag.undefined and not tag.undefined and tag.dimension == existing_tag.dimension
+          if tag != existing_tag and tag in existing_tag.get_derived_bases_list():
+            replace_existing = True
+          if replace_existing:  # Replace the existing by the new tag.
             tags[tags.index(existing_tag)] = tag
             existing_tag = tag
         else:  # no existing tag
@@ -5311,17 +5331,16 @@ def get_common_data(cls, sources, ignore_feature_dim=False, allow_broadcast_all_
     if BehaviorVersion.get() < 11:
       is_equal_opts["broadcast_matches"] = True
     all_dim_tags, tags_dict = Dim.get_all_dimension_tags(sources, is_equal_opts=is_equal_opts)
-    # Check for potential undefined tags, and replace those with defined tags if possible.
-    for axis, dim_tag in enumerate(common.dim_tags):
-      if dim_tag.undefined:
-        other = Dim.get_existing_tag_from_collection(dim_tag, all_dim_tags, is_equal_opts=is_equal_opts)
-        if other and not other.undefined:
-          # We found another dim tag which matches and which is defined.
-          # Replace it.
-          common = common.copy_template_replace_dim_tag(axis=axis, new_dim_tag=other)
     # Check for missing tags, and add those.
     for dim_tag in all_dim_tags:
-      if not Dim.get_existing_tag_from_collection(dim_tag, common.dim_tags, is_equal_opts=is_equal_opts):
+      common_tag = Dim.get_existing_tag_from_collection(dim_tag, common.dim_tags, is_equal_opts=is_equal_opts)
+      if common_tag:
+        # Already have this tag. However, maybe we have a better one.
+        # Dim.get_all_dimension_tags() would have selected that.
+        if dim_tag != common_tag:
+          axis = common.dim_tags.index(common_tag)
+          common = common.copy_template_replace_dim_tag(axis=axis, new_dim_tag=dim_tag)
+      else:
         axis = common.get_default_new_axis_for_dim_tag(dim_tag)
         common = common.copy_add_dim_by_tag(dim_tag, unbroadcast=True, axis=axis)
     if all(s.batch_ndim < common.batch_ndim for s in sources):
@@ -5336,7 +5355,8 @@ def find_matching_dims(self, dim_tag, is_equal_opts):
 
     :param Dim dim_tag:
     :param dict[str,bool]|None is_equal_opts: passed to Dim.is_equal
-    :rtype: list[int] a list of matching axes, counted with batch dim. Sorted in ascending order
+    :rtype: list[int]
+    :return: 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)]
 
@@ -5362,31 +5382,27 @@ def map_other_axis_to_self(other_axis, taken_self_axes):
       :rtype: int
       """
       other_axis_dim_tag = other.get_dim_tag(other_axis)
+      is_equal_opts_ = None
+      matching = None
       # First, try without any is_equal_opts. This is the most restrictive case.
-      is_equal_opts_ = {}
-      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:
-        # Ok, try with the given is_equal_opts.
-        is_equal_opts_ = is_equal_opts.copy()
-        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:
-        # Try harder by allowing broadcasting to match
-        is_equal_opts_["broadcast_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]
-      if not matching:
-        # If still not, then also allow one single dyn_size to be unknown
-        is_equal_opts_["unknown_spatial_matches"] = True
+      # Try with the given is_equal_opts.
+      # Try harder by allowing broadcasting to match.
+      # If still not, then also allow one single dyn_size to be unknown.
+      for opt in [{}, is_equal_opts, "broadcast_matches", "unknown_spatial_matches"]:
+        if isinstance(opt, dict):
+          is_equal_opts_ = opt.copy()
+        elif isinstance(opt, str):
+          if opt in is_equal_opts_:
+            continue
+          is_equal_opts_[opt] = 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. Failing to match axis %s' % (
-          other_axes, other, self, other_axis)
+        if opt == "unknown_spatial_matches":
+          assert len(matching) <= 1, 'cannot match axes %s from %s to %s, failed at other %s, not unique after %s' % (
+            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)
       # If there are multiple matches (e.g. because two axes have the same feature dim), leave their order intact.


Test patch:
diff --git a/tests/test_TFNetworkRecLayer.py b/tests/test_TFNetworkRecLayer.py
--- a/tests/test_TFNetworkRecLayer.py
+++ b/tests/test_TFNetworkRecLayer.py
@@ -4865,6 +4865,17 @@ def test_reclayer_optimize_out_cumsum_step_by_step():
     feat_dim=feat_dim, time_dim=time_dim)
 
 
+def test_reclayer_optimize_out_cumsum_step_by_step_initial():
+  from returnn.tf.util.data import batch_dim, Dim
+  time_dim = SpatialDim("time")
+  feat_dim = FeatureDim("feat", dimension=11)
+  check_reclayer_optimize_out(
+    subnet_layer_dict={
+      "class": "cumsum", "initial_output": 1,
+      "axis": time_dim, "out_shape": {batch_dim, feat_dim}, "out_dim": feat_dim},
+    feat_dim=feat_dim, time_dim=time_dim)
+
+
 def test_reclayer_optimize_out_cumsum_unrelated_axis():
   from returnn.tf.util.data import batch_dim, Dim
   time_dim = SpatialDim("time")
@@ -6125,6 +6136,142 @@ def test_MaskedComputationLayer_UnmaskLayer_in_loop():
         x = y
 
 
+def test_MaskedComputationLayer_UnmaskLayer_in_loop_opt():
+  from test_TFNetworkLayer import make_feed_dict
+  from returnn.tf.layers.basic import SubnetworkLayer
+  from returnn.tf.layers.rec import _SubnetworkRecCell
+  with make_scope() as session:
+    config = Config({"debug_print_layer_output_template": True})
+    net = TFNetwork(
+      extern_data=ExternData({"data": {"dim": 20, "sparse": True}}),
+      config=config)
+    net_dict = {
+      "output": {
+        "class": "rec",
+        "from": "data",
+        "unit": {
+          "const1": {"class": "constant", "value": 1, "with_batch_dim": True},  # just to broadcast mask
+          "mask": {
+            "class": "eval", "from": [":i", "const1"], "out_type": {"dtype": "bool"},
+            "eval": "tf.equal(source(0) % 2, source(1))"},
+          "masked": {
+            "class": "masked_computation", "from": "data:source", "mask": "mask",
+            "unit": {
+              "class": "subnetwork", "from": "data",
+              "subnetwork": {
+                "input0": {"class": "cast", "from": "data", "dtype": "float32"},
+                "input1": {"class": "expand_dims", "axis": "f", "from": "input0"},
+                "output": {"class": "rec", "unit": "cumsum", "n_out": 1, "from": "input1"},
+              },
+            }},
+          "unmask": {"class": "unmask", "from": "masked", "mask": "mask"},
+          "output": {"class": "squeeze", "from": "unmask", "axis": "f"},
+        }
+      }
+    }
+    net.construct_from_dict(net_dict)
+    rec_layer = net.get_layer("output")
+    assert isinstance(rec_layer, RecLayer)
+    rec_cell = rec_layer.cell
+    assert isinstance(rec_cell, _SubnetworkRecCell)
+    assert not rec_cell.layers_in_loop  # all moved out
+    in_data = net.get_layer("data").output
+    out_data = net.get_layer("output").output.copy_as_batch_major()
+    assert in_data.get_time_dim_tag() == out_data.get_time_dim_tag()
+    masked_comp_layer = rec_cell.get_layer_from_outside("masked")
+    assert isinstance(masked_comp_layer, MaskedComputationLayer)
+    masked_comp_sub_layer = masked_comp_layer.sub_layer
+    assert isinstance(masked_comp_sub_layer, SubnetworkLayer)
+    masked_in0_layer = masked_comp_sub_layer.get_sub_layer("input0")
+    masked_in_layer = masked_in0_layer.sources[0]
+    extra = {
+      "mask": rec_cell.get_layer_from_outside("mask").output.placeholder,
+      "masked": masked_comp_layer.output.placeholder,
+      "masked/input0": masked_in0_layer.output.placeholder,
+      "masked/data": masked_in_layer.output.placeholder,
+    }
+    in_v, out_v, out_seq_lens_v, extra_v = session.run(
+      (in_data.placeholder, out_data.placeholder, out_data.get_sequence_lengths(), extra),
+      feed_dict=make_feed_dict(net.extern_data))
+    print(in_v)
+    print(out_v)
+    print("seq lens:", out_seq_lens_v)
+    pprint(extra_v)
+    assert_equal(in_v.shape, out_v.shape)
+    for b in range(in_v.shape[0]):
+      x = 0.0
+      for t in range(in_v.shape[1]):
+        if t >= out_seq_lens_v[b]:
+          continue
+        if t % 2 == 1:
+          y = x + in_v[b, t]
+        else:
+          y = x
+        numpy.testing.assert_almost_equal(y, out_v[b, t])
+        x = y
+
+
+def test_MaskedComputationLayer_in_loop_auto_unmask():
+  # https://github.com/rwth-i6/returnn/issues/769
+  from test_TFNetworkLayer import make_feed_dict
+  from returnn.tf.layers.rec import _SubnetworkRecCell
+  for opt in [False, True]:
+    print("*** using rec optimization:", opt)
+    with make_scope() as session:
+      config = Config({"debug_print_layer_output_template": True})
+      net = TFNetwork(
+        extern_data=ExternData({"data": {"dim": 20, "sparse": True}}),
+        config=config)
+      net_dict = {
+        "output": {
+          "class": "rec",
+          "from": "data",
+          "optimize_move_layers_out": opt,  # test both variants
+          "unit": {
+            "const1": {"class": "constant", "value": 1, "with_batch_dim": True},  # just to broadcast mask
+            "mask": {
+              "class": "eval", "from": [":i", "const1"], "out_type": {"dtype": "bool"},
+              "eval": "tf.equal(source(0) % 2, source(1))"},
+            "in": {"class": "reinterpret_data", "from": "data:source", "set_sparse": False},
+            "masked": {
+              "class": "masked_computation", "from": "in", "mask": "mask",
+              "unit": {"class": "cumsum", "from": "data", "initial_output": 1}},
+            "masked_out": {"class": "copy", "from": "masked"},
+            "output": {"class": "eval", "from": ["masked_out", "in"], "eval": "source(0) + source(1) ** 2"},
+          }
+        }
+      }
+      net.construct_from_dict(net_dict)
+      rec_layer = net.get_layer("output")
+      assert isinstance(rec_layer, RecLayer)
+      rec_cell = rec_layer.cell
+      assert isinstance(rec_cell, _SubnetworkRecCell)
+      if opt:
+        assert not rec_cell.layers_in_loop  # all moved out
+      else:
+        assert not rec_cell.input_layers_moved_out and not rec_cell.output_layers_moved_out  # none moved out
+      in_data = net.get_layer("data").output
+      out_data = net.get_layer("output").output.copy_as_batch_major()
+      print("out:", out_data)
+      assert in_data.get_time_dim_tag() == out_data.get_time_dim_tag()
+      in_v, out_v, out_seq_lens_v = session.run(
+        (in_data.placeholder, out_data.placeholder, out_data.get_sequence_lengths()),
+        feed_dict=make_feed_dict(net.extern_data))
+      print(in_v)
+      print(out_v)
+      print("seq lens:", out_seq_lens_v)
+      assert_equal(in_v.shape, out_v.shape)
+      for b in range(in_v.shape[0]):
+        x = 1
+        for t in range(in_v.shape[1]):
+          if t >= out_seq_lens_v[b]:
+            continue
+          if t % 2 == 1:
+            x = x + in_v[b, t]
+          y = x + in_v[b, t] ** 2
+          numpy.testing.assert_almost_equal(y, out_v[b, t])
+
+
 def test_att_train_search_loss_prev_beam():
   beam_size = 1
   num_ner_labels = 13
diff --git a/tests/test_TFUtil.py b/tests/test_TFUtil.py
--- a/tests/test_TFUtil.py
+++ b/tests/test_TFUtil.py
@@ -401,6 +401,27 @@ def test_Data_get_all_dimension_tags_same_spatial_dim_twice():
   assert all_dim_tags == [batch_dim, time_dim, in_dim, out_dim]
 
 
+def test_Dim_get_all_dimension_tags_one_derived_time():
+  from returnn.tf.util.data import batch_dim
+  time_dim = SpatialDim("time")
+  feat_dim = FeatureDim("feature", 3)
+  a = Data(name="a", dim_tags=[batch_dim, time_dim, feat_dim])
+  print("a:", a)
+  derived_time_dim = SpatialDim("derived_time", derived_from_tag=time_dim)
+  b = Data(name="b", dim_tags=[batch_dim, derived_time_dim, feat_dim])
+  print("b:", b)
+  # Data.get_common_data defaults:
+  is_equal_opts = dict(
+    ignore_feature_dim=False, treat_feature_as_spatial=True,
+    allow_same_spatial_dim=True,
+    undefined_matches=True, derived_matches=True)
+  all_dim_tags, _ = DimensionTag.get_all_dimension_tags([b, a], is_equal_opts=is_equal_opts)
+  print("all_dim_tags:", all_dim_tags)
+  # The is_equal_opts with derived_matches should match time_dim with derived_time_dim.
+  # We explicitly want that the time_dim is returned here.
+  assert derived_time_dim not in all_dim_tags and all_dim_tags == [batch_dim, time_dim, feat_dim]
+
+
 def test_Data_sparse_int32_with_dim_kwargs_init():
   data = Data(name="classes_with_dim", shape=(None,), dim=10, sparse=True, dtype="int32")
   assert data.sparse and data.have_time_axis() and data.shape == (None,) and data.dim == 10