Repository: rwth-i6/returnn

Issue:
Better way to specify axis for layers operating on rec time axis
This is a special case of #597.

I noticed that for #776 in `WindowLayer`, but also for e.g. `CumsumLayer` that operate on the rec time axis, we have a check like this:
```
    if axis == "T" and data.time_dim_axis is None:
      # Assume inside RecLayer.
      # (then the code for a single step)
    else:
      axis = data.get_axis_from_description(axis)
      # (then the code for all steps at once, operating on `axis`)
```

This is flawed as soon as the input has two dynamic dims: Then the input might have a `time_dim_axis` even in the recurrent case, and the layer will silently behave very wrongly.

Instead, I would propose to allow `axis` to be a DimensionTag. Then, we can just check `axis == network.get_rec_parent_layer(inside_loop=False).time_dim_tag`.
Would we also want to allow the `stag:description` syntax here?
Also, in #391 for `CumConcatLayer`, we briefly thought of allowing `:i` as axis name. Maybe allow this too?

I worked on this a month ago, for `CumsumLayer`, see #780. But I never tested this, it was just something a colleague of mine stumbled across but then gave up on. Now with #776, and dim tags becoming popular, I remembered to raise this issue here.


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
@@ -2543,8 +2543,9 @@ def __init__(self, window_size=None, window_dim=None, window_left=None, window_r
       assert window_dim and window_dim.dimension
       window_size = window_dim.dimension
     data = self.input_data.copy_as_batch_major()
-    if axis == "T" and data.time_dim_axis is None:
-      # Assume inside RecLayer.
+    from returnn.tf.util.basic import is_axis_from_description_recurrent
+    if is_axis_from_description_recurrent(axis=axis, network=self.network, data=data):
+      # Inside RecLayer.
       assert self._rec_previous_layer, "%s: expected to be used inside a RecLayer" % self
       assert padding == "same"
       assert window_right is not None or window_left is not None, (
@@ -2590,11 +2591,12 @@ def __init__(self, window_size=None, window_dim=None, window_left=None, window_r
     self.output.placeholder.set_shape(tf.TensorShape(self.output.batch_shape))
 
   @classmethod
-  def get_out_data_from_opts(cls, name, sources, window_size=None, window_dim=None,
+  def get_out_data_from_opts(cls, name, network, sources, window_size=None, window_dim=None,
                              axis="T", out_spatial_dim=None, padding="same", stride=1,
                              **kwargs):
     """
     :param str name:
+    :param returnn.tf.network.TFNetwork network:
     :param list[LayerBase] sources:
     :param int|None window_size:
     :param DimensionTag|None window_dim:
@@ -2610,8 +2612,9 @@ def get_out_data_from_opts(cls, name, sources, window_size=None, window_dim=None
     data = get_concat_sources_data_template(sources)
     data = data.copy_template(name="%s_output" % name)
     data = data.copy_as_batch_major()
-    if axis == "T" and data.time_dim_axis is None:
-      # Assume inside RecLayer.
+    from returnn.tf.util.basic import is_axis_from_description_recurrent
+    if is_axis_from_description_recurrent(axis=axis, network=network, data=data):
+      # Inside RecLayer.
       assert not out_spatial_dim
       new_dim_axis = 1  # after batch
     else:
@@ -2641,9 +2644,10 @@ def get_out_data_from_opts(cls, name, sources, window_size=None, window_dim=None
 
   # noinspection PyMethodOverriding
   @classmethod
-  def get_rec_initial_extra_outputs(cls, batch_dim, rec_layer, window_size=None, window_dim=None,
+  def get_rec_initial_extra_outputs(cls, network, batch_dim, rec_layer, window_size=None, window_dim=None,
                                     axis="T", sources=(), **kwargs):
     """
+    :param returnn.tf.network.TFNetwork network:
     :param tf.Tensor batch_dim:
     :param returnn.tf.layers.rec.RecLayer|LayerBase rec_layer:
     :param int|None window_size:
@@ -2657,8 +2661,9 @@ def get_rec_initial_extra_outputs(cls, batch_dim, rec_layer, window_size=None, w
       window_size = window_dim.dimension
     data = get_concat_sources_data_template(sources)
     data = data.copy_as_batch_major()
-    if axis == "T" and data.time_dim_axis is None:
-      # Assume inside RecLayer.
+    from returnn.tf.util.basic import is_axis_from_description_recurrent
+    if is_axis_from_description_recurrent(axis=axis, network=network, data=data):
+      # Inside RecLayer.
       shape = list(data.batch_shape)
       shape[0] = batch_dim
       shape.insert(1, window_size)
@@ -2684,9 +2689,8 @@ def __init__(self, axis="T", additional_left_summand_per_element=None, reverse=F
     x = data.placeholder
     if additional_left_summand_per_element is not None:
       x = additional_left_summand_per_element + x
-    if axis == "T" and data.time_dim_axis is None:
-      # Assume inside RecLayer.
-      assert self._rec_previous_layer, "%s: expected to be used inside a RecLayer" % self
+    from returnn.tf.util.basic import is_axis_from_description_recurrent
+    if is_axis_from_description_recurrent(axis=axis, network=self.network, data=data):
       assert not reverse
       prev_state = self._rec_previous_layer.rec_vars_outputs["state"]
       next_state = prev_state + x
@@ -2710,9 +2714,11 @@ def get_out_data_from_opts(cls, name, sources, axis="T", **kwargs):
     # Just same format.
     return get_concat_sources_data_template(sources, name="%s_output" % name)
 
+  # noinspection PyMethodOverriding
   @classmethod
-  def get_rec_initial_extra_outputs(cls, batch_dim, rec_layer, axis="T", sources=(), **kwargs):
+  def get_rec_initial_extra_outputs(cls, network, batch_dim, rec_layer, axis="T", sources=(), **kwargs):
     """
+    :param returnn.tf.network.TFNetwork network:
     :param tf.Tensor batch_dim:
     :param returnn.tf.layers.rec.RecLayer|LayerBase rec_layer:
     :param str axis:
@@ -2720,8 +2726,8 @@ def get_rec_initial_extra_outputs(cls, batch_dim, rec_layer, axis="T", sources=(
     :rtype: dict[str,tf.Tensor]
     """
     data = get_concat_sources_data_template(sources)
-    if axis == "T" and data.time_dim_axis is None:
-      # Assume inside RecLayer.
+    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)}
     return {}
diff --git a/returnn/tf/util/basic.py b/returnn/tf/util/basic.py
--- a/returnn/tf/util/basic.py
+++ b/returnn/tf/util/basic.py
@@ -6976,3 +6976,19 @@ def _callback(self, value):
     self.most_recent_value = value
     self.callback_count += 1
     return 0
+
+
+def is_axis_from_description_recurrent(axis, network, data):
+  """
+  :param str|DimensionTag axis:
+  :param returnn.tf.network.TFNetwork network:
+  :param Data data:
+  :rtype: bool
+  """
+  if not network.is_inside_rec_layer(inside_loop=True):
+    return False
+  if isinstance(axis, str) and axis.lower() == "t" and data.time_dim_axis is None:  # legacy case
+    return True
+  if isinstance(axis, DimensionTag) and axis == network.get_inside_rec_time_dim(inside_loop=True):
+    return True
+  return False


Test patch:
diff --git a/tests/test_TFNetworkRecLayer.py b/tests/test_TFNetworkRecLayer.py
--- a/tests/test_TFNetworkRecLayer.py
+++ b/tests/test_TFNetworkRecLayer.py
@@ -3436,6 +3436,8 @@ def check_reclayer_optimize_out(subnet_layer_dict, other_subnet_layers=None, sha
   if not time_dim:
     time_dim = DimensionTag(kind=DimensionTag.Types.Spatial, description="time")
   n_out = subnet_layer_dict.get("n_out", 17)
+  if subnet_layer_dict.get("out_dim", None):
+    n_out = subnet_layer_dict["out_dim"].dimension
   n_batch = 5
   n_time = 7
   subnet_layer_dict["n_out"] = n_out
@@ -4609,6 +4611,26 @@ def make_extern_data():
     assert "encoder_int" in cell.input_layers_moved_out
 
 
+def test_reclayer_optimize_out_cumsum_step_by_step():
+  from returnn.tf.util.data import BatchDim, DimensionTag
+  time_dim = DimensionTag(kind=DimensionTag.Types.Spatial, description="time")
+  feat_dim = DimensionTag(kind=DimensionTag.Types.Feature, description="feat", dimension=11)
+  check_reclayer_optimize_out(
+    subnet_layer_dict={
+      "class": "cumsum", "axis": time_dim, "out_shape": {BatchDim, 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 BatchDim, DimensionTag
+  time_dim = DimensionTag(kind=DimensionTag.Types.Spatial, description="time")
+  feat_dim = DimensionTag(kind=DimensionTag.Types.Feature, description="feat", dimension=11)
+  check_reclayer_optimize_out(
+    subnet_layer_dict={
+      "class": "cumsum", "axis": feat_dim, "out_shape": {BatchDim, feat_dim}, "out_dim": feat_dim},
+    feat_dim=feat_dim, time_dim=time_dim)
+
+
 def test_subnet_load_on_init_rec():
   import tempfile
   model_tmp_dir = tempfile.mkdtemp("tmp-checkpoint")