Repository: rwth-i6/returnn

Issue:
Layer axes or axis option should not allow int
This requires a new behavior version (#508).



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
@@ -1862,7 +1862,7 @@ def __init__(self, axis=None, energy_factor=None,
                start=None, window_start=None, window_size=None, use_time_mask=None,
                log_space=False, **kwargs):
     """
-    :param str|None axis: which axis to do the softmax over
+    :param DimensionTag|str|None axis: which axis to do the softmax over. "T" by default
     :param float|None energy_factor: the energy will be scaled by this factor.
       This is like a temperature for the softmax.
       In Attention-is-all-you-need, this is set to 1/sqrt(base_ctx.dim).
@@ -1895,7 +1895,7 @@ def __init__(self, axis=None, energy_factor=None,
     if use_time_mask:
       energy_mask = SeqLenMaskLayer.build_mask(
         energy_data,
-        axis=axis,
+        axis=axis, axis_allow_int=True,
         start=start.output if start else None,
         window_start=window_start.output if isinstance(window_start, LayerBase) else window_start,
         window_size=window_size.output if isinstance(window_size, LayerBase) else window_size)
@@ -1927,7 +1927,7 @@ def get_dep_layers(self):
   def _get_axis_to_reduce(cls, input_data, axis, exception_prefix):
     """
     :param Data input_data:
-    :param str|None axis:
+    :param DimensionTag|str|None axis:
     :param str|object exception_prefix:
     :rtype: int
     """
@@ -1943,7 +1943,7 @@ def get_out_data_from_opts(cls, name, sources, axis=None, start=None, window_sta
     """
     :param str name:
     :param list[LayerBase] sources:
-    :param str|None axis:
+    :param DimensionTag|str|None axis:
     :param LayerBase|None start:
     :param LayerBase|None window_start:
     :param LayerBase|int|None window_size:
@@ -2018,10 +2018,14 @@ def __init__(self, mask_value, axis="T", seq_len_source=None,
       self.allow_inf_in_output = True
 
   @classmethod
-  def build_mask(cls, x, axis="T", seq_len_source=None, start=None, window_start=None, window_size=None):
+  def build_mask(cls, x, axis="T", axis_allow_int=NotSpecified,
+                 seq_len_source=None, start=None, window_start=None, window_size=None):
     """
     :param Data x:
-    :param str|int axis:
+    :param DimensionTag|str|int axis:
+    :param bool|NotSpecified axis_allow_int:
+      Some callers of this function would pass in an int for axis directly.
+      In that case, explicitly set this to True.
     :param Data|None seq_len_source:
     :param Data|None start:
     :param Data|None window_start:
@@ -2032,7 +2036,7 @@ def build_mask(cls, x, axis="T", seq_len_source=None, start=None, window_start=N
     from returnn.tf.util.basic import get_shape
     energy = x.placeholder
     energy_shape = get_shape(energy)
-    axis = x.get_axis_from_description(axis)
+    axis = x.get_axis_from_description(axis, allow_int=axis_allow_int)
     assert x.is_axis_dynamic(axis), "%s: use_time_mask True, dyn time axis expected" % x
     if seq_len_source:
       energy_mask = seq_len_source.copy_compatible_to(x).get_sequence_mask_broadcast(axis=axis)
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
@@ -3415,13 +3415,43 @@ def get_axes(self, exclude_time=False, exclude_batch=False, exclude_feature=Fals
       axes.pop(axes.index(self.feature_dim_axis))
     return axes
 
-  def get_axes_from_description(self, axes, allow_int=True):
+  @classmethod
+  def _verify_axis_int_from_description(cls, allow_int=NotSpecified):
+    """
+    Call this when you have the case that ``axis`` or ``axes``
+    in :func:`get_axes_from_description` or :func:`get_axis_from_description`
+    was specified as int.
+
+    :param bool|NotSpecified allow_int:
+    """
+    msg = "Do not specify axis as int but as str or DimensionTag instead."
+    if allow_int is NotSpecified:
+      from returnn.util import BehaviorVersion
+      BehaviorVersion.require(condition=False, message=msg, version=5)
+    if allow_int:
+      return
+    raise Exception(msg)
+
+  def _make_valid_int_axis(self, axis):
+    """
+    :param int axis: counted with batch. anything in [-ndim,ndim-1]
+    :return: axis in [0,ndim-1]
+    :rtype: int
+    """
+    if axis < 0:
+      assert axis + self.batch_ndim >= 0
+      axis += self.batch_ndim
+    assert axis < self.batch_ndim
+    return axis
+
+  def get_axes_from_description(self, axes, allow_int=NotSpecified):
     """
     :param int|list[int]|str|list[str|DimensionTag]|DimensionTag|None axes: one axis or multiple axis, or none.
       This is counted with batch-dim, which by default is axis 0 (see enforce_batch_dim_axis).
       It also accepts the special tokens "B"|"batch", "spatial", "spatial_except_time", or "F"|"feature",
       and more (see the code).
-    :param bool allow_int: whether to allow an int directly. in almost all cases, it is better to use a symbolic name
+    :param bool|NotSpecified allow_int: whether to allow an int directly.
+      in almost all cases, it is better to use a symbolic name
       to specify an axis, as different layers could reorder them, and maybe also change their behavior in the future.
     :return: list of axes, counted with batch-dim
     :rtype: list[int]
@@ -3429,55 +3459,57 @@ def get_axes_from_description(self, axes, allow_int=True):
     if axes is None or axes == "":
       return []
     if isinstance(axes, DimensionTag):
+      # Once we have not guaranteed unique dim tags, multiple axes could match.
+      # https://github.com/rwth-i6/returnn/issues/632
       return [i for (i, tag) in enumerate(self.dim_tags) if tag == axes]
-    if not allow_int:
-      assert not isinstance(axes, int)
+    if isinstance(axes, int):
+      self._verify_axis_int_from_description(allow_int=allow_int)
+      return [self._make_valid_int_axis(axes)]
     assert isinstance(axes, (str, int, list, tuple))
-    if isinstance(axes, (list, tuple)):
-      assert all([a is None or isinstance(a, (str, int, DimensionTag)) for a in axes])
-      if not allow_int:
-        assert all([not isinstance(a, int) for a in axes])
     if isinstance(axes, str):
       import re
       axes = axes.lower()
       if axes in ["b", "batch"]:
         assert self.batch_dim_axis is not None
-        axes = self.batch_dim_axis
+        return [self.batch_dim_axis]
       elif axes == "spatial":
-        axes = self.get_spatial_batch_axes()
+        return self.get_spatial_batch_axes()
       elif re.match("(s|spatial):-?\\d+$", axes):
         s = int(axes.split(":")[1])
         spatial_axes = self.get_spatial_batch_axes()
         if s < 0:
           s += len(spatial_axes)
         assert s < len(spatial_axes), "%s get_axes_from_description: %r invalid" % (self, axes)
-        axes = spatial_axes[s]
+        return [spatial_axes[s]]
       elif axes in ["dyn", "dynamic"]:
-        axes = self.get_dynamic_axes()
+        return self.get_dynamic_axes()
       elif re.match("(d|dyn|dynamic):-?\\d+$", axes):
         s = int(axes.split(":")[1])
         dyn_axes = self.get_dynamic_axes()
         if s < 0:
           s += len(dyn_axes)
         assert 0 <= s < len(dyn_axes), "%s get_axes_from_description: %r invalid" % (self, axes)
-        axes = dyn_axes[s]
+        return [dyn_axes[s]]
       elif axes == "spatial_except_time":
         axes = self.get_spatial_batch_axes()
         assert self.time_dim_axis is not None
         axes.remove(self.time_dim_axis)
+        return axes
       elif axes in ["t", "time"]:
         assert self.time_dim_axis is not None
-        axes = self.time_dim_axis
+        return [self.time_dim_axis]
       elif axes == "t?":
-        axes = [self.time_dim_axis] if self.time_dim_axis is not None else []
+        return [self.time_dim_axis] if self.time_dim_axis is not None else []
       elif axes == "except_time":  # also except batch
         axes = list(range(self.batch_ndim))
         axes.remove(self.batch_dim_axis)
         if self.time_dim_axis is not None:
           axes.remove(self.time_dim_axis)
+        return axes
       elif axes == "except_batch":
         axes = list(range(self.batch_ndim))
         axes.remove(self.batch_dim_axis)
+        return axes
       elif re.match("(except_batch):-?\\d+$", axes):
         s = int(axes.split(":")[1])
         non_batch_axes = list(range(self.batch_ndim))
@@ -3486,43 +3518,40 @@ def get_axes_from_description(self, axes, allow_int=True):
         if s < 0:
           s += len(non_batch_axes)
         assert 0 <= s < len(non_batch_axes), "%s get_axes_from_description: %r invalid" % (self, axes)
-        axes = non_batch_axes[s]
+        return [non_batch_axes[s]]
       elif axes == "*":
-        axes = list(range(self.batch_ndim))
+        return list(range(self.batch_ndim))
       elif axes == "static":
-        axes = self.get_static_axes()
+        return self.get_static_axes()
       elif re.match("(static):-?\\d+$", axes):
         s = int(axes.split(":")[1])
         static_axes = self.get_static_axes()
         if s < 0:
           s += len(static_axes)
         assert 0 <= s < len(static_axes), "%s get_axes_from_description: %r invalid" % (self, axes)
-        axes = static_axes[s]
+        return [static_axes[s]]
       elif axes in ["f", "feature", "non_spatial"]:
-        axes = self.get_feature_batch_axes()
+        return self.get_feature_batch_axes()
       elif all([a in "btf" for a in axes]):
         return self.get_axes_from_description(list(axes))
       elif axes.startswith("stag:"):  # spatial tag
-        axes = self.get_axis_by_tag_name(axes[len("stag:"):], spatial_only=True)
+        return [self.get_axis_by_tag_name(axes[len("stag:"):], spatial_only=True)]
       elif axes.startswith("stag-single:"):  # spatial tag which possibly matches multiple spatial axes
         # in this case, a name of form "stag-single:<idx>:<name> is expected.
         # idx is relative to the matching stags, i.e., it is the index among the list of spatial dims matching the name
         _, idx_s, name = axes.split(":", 2)  # stag-single:<idx>:<name>
         idx = int(idx_s)
-        axes = self.get_axes_by_tag_name(name, spatial_only=True)[idx]
-      else:
-        raise Exception("invalid axis mode %r" % axes)
-    if isinstance(axes, int):
-      axes = [axes]
-    assert isinstance(axes, (tuple, list)), "invalid axis %r" % axes
+        return [self.get_axes_by_tag_name(name, spatial_only=True)[idx]]
+      raise Exception("invalid axis mode %r" % axes)
+    assert isinstance(axes, (tuple, list)), "invalid axes %r" % axes
     flat_axes = []
     for i in axes:
       if isinstance(i, int):
-        flat_axes += [i]
+        self._verify_axis_int_from_description(allow_int=allow_int)
+        flat_axes.append(self._make_valid_int_axis(i))
       else:
         assert isinstance(i, (str, tuple, list, DimensionTag))
-        flat_axes += self.get_axes_from_description(i)
-    flat_axes = [i % self.batch_ndim for i in flat_axes]
+        flat_axes += self.get_axes_from_description(i, allow_int=allow_int)
     res = []
     for i in flat_axes:
       if i not in res:
@@ -3538,10 +3567,10 @@ def get_dim_tag_from_description(self, axis):
     axis_int = self.get_axis_from_description(axis, allow_int=False)
     return self.dim_tags[axis_int]
 
-  def get_axis_from_description(self, axis, allow_int=True):
+  def get_axis_from_description(self, axis, allow_int=NotSpecified):
     """
     :param int|str|DimensionTag axis:
-    :param bool allow_int:
+    :param bool|NotSpecified allow_int:
     :return: axis, counted with batch-dim
     :rtype: int
     """
diff --git a/returnn/util/basic.py b/returnn/util/basic.py
--- a/returnn/util/basic.py
+++ b/returnn/util/basic.py
@@ -209,7 +209,7 @@ class BehaviorVersion:
   The version will be set after the config is defined at __main__.init_config() or Engine.__init__()
   """
 
-  _latest_behavior_version = 4
+  _latest_behavior_version = 5
   _behavior_version = None  # type: typing.Optional[int]
 
   @classmethod


Test patch:
diff --git a/tests/test_TFNetworkLayer.py b/tests/test_TFNetworkLayer.py
--- a/tests/test_TFNetworkLayer.py
+++ b/tests/test_TFNetworkLayer.py
@@ -2421,7 +2421,7 @@ def test_ScatterNdLayer_RangeLayer():
     "range": {"class": "range", "limit": n_ts, "out_spatial_dim": ts_dim},  # (Ts,)
     "add_t": {"class": "combine", "kind": "add", "from": ["t", "range"], "out_shape": {BatchDim, ts_dim}},  # (B,Ts)
     "t_rel_var": {"class": "variable", "shape": (n_ts, n_out), "init": "glorot_uniform"},  # (B,Ts,D)
-    "output": {"class": "scatter_nd", "from": "t_rel_var", "position": "add_t", "position_axis": -1,
+    "output": {"class": "scatter_nd", "from": "t_rel_var", "position": "add_t", "position_axis": ts_dim,
                "output_dim_via_time_from": "data", "filter_invalid_indices": True}
   }
   with make_scope() as session:
diff --git a/tests/test_TFNetworkRecLayer.py b/tests/test_TFNetworkRecLayer.py
--- a/tests/test_TFNetworkRecLayer.py
+++ b/tests/test_TFNetworkRecLayer.py
@@ -5109,11 +5109,12 @@ def add_lstm(i, direction, src):
       network[name] = {"class": "rec", "unit": "lstmp", "n_out": lstm_dim, "dropout": 0.1, "L2": 0.01, "direction": 1,
                        "from": src}
       return name
-    network["%s_win" % name] = {"class": "window", "window_size": lstm_window, "window_right": lstm_window - 1,
-                                "from": src}  # (B,T,W,D)
+    network["%s_win" % name] = {
+      "class": "window", "window_size": lstm_window, "window_right": lstm_window - 1, "from": src}  # (B,T,W,D)
     network["%s_mdims" % name] = {"class": "merge_dims", "axes": "BT", "from": ["%s_win" % name]}  # (B*T,W,D)
-    network["%s_rdims" % name] = {"class": "reinterpret_data", "enforce_batch_major": True, "set_axes": {"T": 1},
-                                  "from": ["%s_mdims" % name]}  # (B*T,W,D)
+    network["%s_rdims" % name] = {
+      "class": "reinterpret_data", "enforce_batch_major": True, "set_axes": {"T": "spatial"},
+      "from": ["%s_mdims" % name]}  # (B*T,W,D)
     network["%s_rec" % name] = {
       "class": "rec", "unit": "lstmp", "n_out": lstm_dim, "dropout": 0.1, "L2": 0.01, "direction": -1,
       "from": ["%s_rdims" % name]}  # (B*T,W,D')