Repository: rwth-i6/returnn

Issue:
`RecUnstackLayer` side effect of marking axis the same can be confusing
The use of `mark_same_time`.

I think it is a better solution when the user explicitly specifies the axis to `RecLayer`.

Or at least this behavior should be an option to `RecUnstackLayer` which should be disabled by default. So that it is explicit for the user.


Discussion:

Resolving patch:
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
@@ -98,7 +98,7 @@ def __init__(self,
                use_global_rec_step_offset=False,
                include_eos=False,
                debug=None,
-               _time_dim_tag=None,
+               axis=None,
                **kwargs):
     """
     :param str|_SubnetworkRecCell unit: the RNNCell/etc name, e.g. "nativelstm". see comment below.
@@ -120,7 +120,7 @@ def __init__(self,
     :param bool use_global_rec_step_offset:
     :param bool include_eos: for search, whether we should include the frame where "end" is True
     :param bool|None debug:
-    :param DimensionTag|None _time_dim_tag:
+    :param DimensionTag|None axis:
     """
     super(RecLayer, self).__init__(**kwargs)
     import re
@@ -147,7 +147,7 @@ def __init__(self,
     self._initial_state_deps = [layer for layer in nest.flatten(initial_state) if isinstance(layer, LayerBase)]
     self._input_projection = input_projection
     self._max_seq_len = max_seq_len
-    self.time_dim_tag = _time_dim_tag
+    self.time_dim_tag = axis
     self.include_eos = include_eos
     if optimize_move_layers_out is None:
       optimize_move_layers_out = self.network.get_config().bool("optimize_move_layers_out", True)
@@ -341,7 +341,11 @@ def transform_config_dict(cls, d, network, get_layer):
         time_dim_tag = DimensionTag(
           description="dyn-time:%s%s" % (network.get_absolute_name_prefix(), d["_name"]),
           kind=DimensionTag.Types.Time)
-    d["_time_dim_tag"] = time_dim_tag
+    if d.get("axis") and time_dim_tag:
+      time_dim_tag_explicit = d["axis"]
+      assert isinstance(time_dim_tag_explicit, DimensionTag)
+      time_dim_tag.declare_same_as(time_dim_tag_explicit)
+    d["axis"] = time_dim_tag
 
     if isinstance(d.get("unit"), dict):
       sub_net_dict = d.pop("unit")
@@ -382,11 +386,11 @@ def max_len_from(src):
       d["max_seq_len"] = eval(d["max_seq_len"], {"max_len_from": max_len_from, "tf": tf})
 
   @classmethod
-  def get_out_data_from_opts(cls, network, unit, _time_dim_tag=None, sources=(), initial_state=None, **kwargs):
+  def get_out_data_from_opts(cls, network, unit, axis=None, sources=(), initial_state=None, **kwargs):
     """
     :param returnn.tf.network.TFNetwork network:
     :param str|dict[str] unit:
-    :param DimensionTag|None _time_dim_tag:
+    :param DimensionTag|None axis:
     :param list[LayerBase] sources:
     :param str|LayerBase|list[str|LayerBase] initial_state:
     :rtype: Data
@@ -419,8 +423,8 @@ def get_out_data_from_opts(cls, network, unit, _time_dim_tag=None, sources=(), i
         out = out.copy_as_time_batch_major()  # Otherwise the output is always [T,B,F]
     else:
       raise Exception("n_out or out_type must be specified")
-    if out.have_time_axis() and _time_dim_tag:
-      out = out.copy_template_replace_dim_tag(axis=out.time_dim_axis, new_dim_tag=_time_dim_tag)
+    if out.have_time_axis() and axis:
+      out = out.copy_template_replace_dim_tag(axis=out.time_dim_axis, new_dim_tag=axis)
     for dep in deps:
       if dep:
         out.beam = SearchBeam.get_combined_beam(out.beam, dep.output.beam)
@@ -4604,15 +4608,22 @@ class RecUnstackLayer(LayerBase):
   (See get_input_moved_out for some internal handling.)
 
   Effectively, this layer is very similar to :class:`CopyLayer`,
-  with the only special behavior that it assigns the loop dimension of RecLayer.
+  with the only special behavior that it checks (or even assigns) the loop dimension of RecLayer.
   """
   layer_class = "rec_unstack"
 
-  def __init__(self, axis, **kwargs):
+  def __init__(self, axis=None, declare_rec_time=False, **kwargs):
     """
-    :param str|DimensionTag axis:
+    Due to automatic optimization, not much happens here.
+    The real logic happens in :func:`get_out_data_from_opts`.
+
+    Note that it is allowed to leave both `axis` and `declare_rec_time` unset,
+    in case you assign `axis` to the rec layer, and the source here has the same axis (dim tag).
+
+    :param str|DimensionTag|None axis:
+    :param bool declare_rec_time:
     """
-    axis  # noqa  # unused here, used in get_out_data_from_opts
+    axis, declare_rec_time  # noqa  # unused here, used in get_out_data_from_opts
     super(RecUnstackLayer, self).__init__(**kwargs)
     assert len(self.sources) == 1
     src = self.sources[0].output
@@ -4622,22 +4633,31 @@ def __init__(self, axis, **kwargs):
     self.output.placeholder = src.placeholder
 
   @classmethod
-  def get_out_data_from_opts(cls, name, axis, sources, network, **kwargs):
+  def get_out_data_from_opts(cls, name, sources, network, axis=None, declare_rec_time=False, **kwargs):
     """
     :param str name:
-    :param str|DimensionTag axis:
     :param list[LayerBase] sources:
     :param returnn.tf.network.TFNetwork network:
+    :param str|DimensionTag|None axis:
+    :param bool declare_rec_time:
     :rtype: Data
     """
     assert sources
     out = sources[0].output.copy_template(name="%s_output" % name)
+    if not axis:
+      axis = network.get_inside_rec_time_dim(inside_loop=False)
+      assert axis, "%s %r: expected to be inside rec layer" % (cls.__name__, name)
     out_dim = out.get_dim_tag_from_description(axis)
     rec_time_dim = network.get_inside_rec_time_dim()
     if rec_time_dim:
       if rec_time_dim.is_dim_known():  # defined
-        assert out_dim == rec_time_dim
+        assert out_dim == rec_time_dim, "%s %r: rec dim %s does not match input %s dim %s" % (
+          cls.__name__, name, rec_time_dim, sources[0], out_dim)
       else:
+        assert out_dim.is_dim_known()  # it comes from the input, so it should be known
+        assert out_dim != rec_time_dim  # rec_time_dim is unknown, so it cannot be the same
+        assert declare_rec_time, "%s %r: must either set known axis on rec %s or enable declare_rec_time" % (
+          cls.__name__, name, rec_time_dim)
         rec_time_dim.declare_same_as(out_dim)
       out.mark_same_time(out_dim, must_match=True)
       return out.copy_template_excluding_time_dim()


Test patch:
diff --git a/tests/test_TFNetworkRecLayer.py b/tests/test_TFNetworkRecLayer.py
--- a/tests/test_TFNetworkRecLayer.py
+++ b/tests/test_TFNetworkRecLayer.py
@@ -1183,10 +1183,38 @@ def test_rec_explicit_lstm():
     pprint(res)
 
 
-def test_RecUnstackLayer_rec_no_input():
+def test_RecUnstackLayer_rec_no_input_explicit_axis():
+  time_dim = DimensionTag(kind=DimensionTag.Types.Spatial, description="time")
+  net_dict = {
+    "output": {"class": "rec", "from": [], "axis": time_dim, "unit": {
+      "input": {"class": "rec_unstack", "from": "base:data:data"},
+      "output": {"class": "combine", "kind": "add", "from": ["prev:output", "input"]},
+    }},
+  }
+  config = Config({
+    "extern_data": {"data": {"dim": 5, "same_dim_tags_as": {"T": time_dim}}},
+  })
+  with make_scope() as session:
+    net = TFNetwork(config=config)
+    net.construct_from_dict(net_dict)
+    rec_layer = net.get_layer("output")
+    assert isinstance(rec_layer, RecLayer)
+    cell = rec_layer.cell
+    from returnn.tf.layers.rec import _SubnetworkRecCell
+    assert isinstance(cell, _SubnetworkRecCell)
+    assert_equal(cell.input_layers_moved_out, ["input"])
+    assert_equal(cell.layers_in_loop, ["output"])
+    in_data = net.extern_data.get_default_input_data()
+    out_data = rec_layer.output
+    assert_equal(in_data.get_time_dim_tag(), out_data.get_time_dim_tag())
+    from test_TFNetworkLayer import make_feed_dict
+    session.run(out_data.placeholder, feed_dict=make_feed_dict([in_data]))
+
+
+def test_RecUnstackLayer_rec_no_input_declare_rec_time():
   net_dict = {
     "output": {"class": "rec", "from": [], "unit": {
-      "input": {"class": "rec_unstack", "from": "base:data:data", "axis": "T"},
+      "input": {"class": "rec_unstack", "from": "base:data:data", "axis": "T", "declare_rec_time": True},
       "output": {"class": "combine", "kind": "add", "from": ["prev:output", "input"]},
     }},
   }