Repository: rwth-i6/returnn

Issue:
Tensor.feature_dim and Tensor.feature_dim_axis
The question is whether to use them in new code, esp whether it is used in the new frontend API (https://github.com/rwth-i6/returnn/issues/1264, https://github.com/rwth-i6/returnn/issues/1120), or whether they should be avoided. See also here for some initial comments: https://github.com/rwth-i6/returnn_common/issues/252#issuecomment-1460933674


Another related question: Should it be like in REUTURNN-common, an alias to `feature_dim_or_sparse_dim`? Or just when it is not sparse?
 **Edit** `feature_dim` is only about the dense dim (via https://github.com/rwth-i6/returnn_common/pull/255). There is `sparse_dim` for the sparse dim.

Some arguments to use `feature_dim`:

- The code might be too verbose otherwise. It potentially can simplify a lot of code.
- There are already a lot of usages via returnn-common `nn`. When the frontend API becomes the underlying core behind returnn-common `nn`, it would break a lot of code.

Some arguments against `feature_dim`:

- There might be ambiguous cases, where it is not clear what `feature_dim` is set to. We should try to make sure such cases do not exist, and `feature_dim` is just None after any ambiguous op, like `split_dims`, or when the dim was removed, like `reduce`. But for example, what about `dot`? Would `dot` handle it, or is it up to `Linear`? Similarly `conv`, `lstm`, etc.
- There might be code which works if `feature_dim` is set to one dim and then breaks when it is set otherwise, and this is maybe difficult to debug. It might make it harder to reason about some code. It adds another thing to think about.
- It can make some internal code more complicated. E.g. inside of `dot`, we need some extra handling to correctly set the returned `feature_dim`.

Let's collect some more arguments, towards either side, to really get a complete picture here. Maybe also let's post some exemplary positive or negative examples.


Discussion:
Currently, `feature_dim` is internally used (as default axis, or just always) in:

* `cross_entropy` and other loss functions. CE has explicit code for both sparse and dense, reduces to `-nn.dot(target, log_prob, reduce=axis)` in the generic dense case, `nn.gather` for sparse.
* `ctc_greedy_decode`. expects dense.
* `label_smoothing`. explicit code for both sparse and dense cases.
* `concat_features`. expects dense.
* `gating`. expects dense.
* `Transformer`, `PrevConvNet`, some others. mostly expect dense.


Resolving patch:
diff --git a/returnn/tensor/_tensor_extra.py b/returnn/tensor/_tensor_extra.py
--- a/returnn/tensor/_tensor_extra.py
+++ b/returnn/tensor/_tensor_extra.py
@@ -27,7 +27,6 @@ def __init__(
         *,
         tensor: Tensor,
         time_dim_axis=NotSpecified,
-        feature_dim_axis=NotSpecified,
         available_for_inference=True,
         batch=None,
         beam=None,
@@ -37,7 +36,6 @@ def __init__(
         :param tensor:
         :param int|None|NotSpecified time_dim_axis: where we have the time dim axis, after we added the batch-dim.
             this is often 1. however, can be None if there is no time-dim.
-        :param int|None|NotSpecified feature_dim_axis: feature dim axis. by default, it's the last one
         :param bool available_for_inference: e.g. the extern data "classes" is usually not available for inference
         :param BatchInfo|None batch:
         :param SearchBeam|None beam: the batch-dim could be extended by a beam-size,
@@ -54,21 +52,10 @@ def __init__(
         self.control_flow_ctx = control_flow_ctx
         self.available_for_inference = available_for_inference
 
-        if feature_dim_axis is NotSpecified:
-            pass
-        elif feature_dim_axis is None:
-            pass
-        elif isinstance(feature_dim_axis, int):
-            assert self.tensor.version == 1
-            assert not self.tensor.sparse, "cannot have feature_dim_axis when sparse"
-            if feature_dim_axis < 0:
-                feature_dim_axis += self.tensor.batch_ndim
-            assert 0 <= feature_dim_axis < self.tensor.batch_ndim
-        else:
-            raise TypeError(f"unexpected feature_dim_axis type {type(feature_dim_axis)}")
-        self.feature_dim_axis = feature_dim_axis
         if time_dim_axis is NotSpecified:
-            pass
+            # just like feature_dim_axis, time_dim_axis is just None with version>=2, without any default behavior.
+            if self.tensor.version >= 2:
+                time_dim_axis = None
         elif time_dim_axis is None:
             pass
         elif isinstance(time_dim_axis, int):
@@ -89,6 +76,7 @@ class _TensorMixin:
     _dims: Tuple[Dim, ...]
     dtype: str
     sparse_dim: Optional[Dim]
+    _feature_dim_axis: Optional[Union[int, NotSpecified]]
     _raw_tensor: Optional[_t.RawTensorType]
     raw_tensor: Optional[_t.RawTensorType]
     version: int
@@ -1169,6 +1157,9 @@ def copy_compatible_to(
                 v.feature_dim_axis = NotSpecified
                 if v.feature_dim_axis != data.feature_dim_axis:
                     v.feature_dim_axis = data.feature_dim_axis
+        else:
+            if v.feature_dim_axis != data.feature_dim_axis:
+                v.feature_dim_axis = data.feature_dim_axis
 
         # Reset sparse
         if self.sparse:
@@ -1807,10 +1798,8 @@ def feature_dim_axis(self):
         :return: feature dim axis, counted with batch-dim
         :rtype: int|None
         """
-        if self.version >= 2:
-            return None
-        if self.feature_dim_axis_or_unspecified is not NotSpecified:
-            return self.feature_dim_axis_or_unspecified
+        if self._feature_dim_axis is not NotSpecified:
+            return self._feature_dim_axis
         return self._default_feature_dim_axis()
 
     @feature_dim_axis.setter
@@ -1821,12 +1810,11 @@ def feature_dim_axis(self: _t.Tensor, value):
         assert value is NotSpecified or value is None or isinstance(value, int)
         if self.feature_dim_axis_or_unspecified == value:
             return
-        if self.version >= 2 and value in (None, NotSpecified):
-            return
-        assert self.version == 1, "feature_dim_axis is deprecated"
+        if self.version >= 2 and value is NotSpecified:
+            value = None
         if isinstance(value, int):
             assert 0 <= value < self.batch_ndim
-        self._make_extra().feature_dim_axis = value
+        self._feature_dim_axis = value
 
     @property
     def feature_dim_axis_or_unspecified(self):
@@ -1834,11 +1822,7 @@ def feature_dim_axis_or_unspecified(self):
         :return: feature dim axis, counted with batch-dim. could also be unspecified
         :rtype: int|None|NotSpecified
         """
-        if self.version >= 2:
-            return None
-        if not self._extra:
-            return NotSpecified
-        return self._extra.feature_dim_axis
+        return self._feature_dim_axis
 
     @property
     def time_dim_axis(self) -> Optional[int]:
diff --git a/returnn/tensor/tensor.py b/returnn/tensor/tensor.py
--- a/returnn/tensor/tensor.py
+++ b/returnn/tensor/tensor.py
@@ -21,8 +21,9 @@
 """
 
 from __future__ import annotations
-from typing import Optional, Sequence, Tuple, Generic, TypeVar
+from typing import Optional, Union, Sequence, Tuple, Generic, TypeVar
 
+from returnn.util.basic import NotSpecified
 from .dim import Dim
 from ._tensor_extra import _TensorExtra, _TensorMixin
 from . import _tensor_extra
@@ -37,12 +38,13 @@ class Tensor(_TensorMixin, Generic[RawTensorType]):
 
     size_dtype = "int32"
 
-    __slots__ = ("name", "_dims", "dtype", "sparse_dim", "_raw_tensor", "version", "_extra")
+    __slots__ = ("name", "_dims", "dtype", "sparse_dim", "_feature_dim_axis", "_raw_tensor", "version", "_extra")
 
     name: str
     _dims: Tuple[Dim, ...]
     dtype: str
     sparse_dim: Optional[Dim]
+    _feature_dim_axis: Optional[Union[int, NotSpecified]]
     _raw_tensor: Optional[RawTensorType]
     version: int
     _extra: Optional[_TensorExtra]
@@ -54,6 +56,7 @@ def __init__(
         dtype: Optional[str] = None,
         *,
         sparse_dim: Optional[Dim] = None,
+        feature_dim_axis: Optional[Union[int, NotSpecified]] = NotSpecified,
         raw_tensor: Optional[RawTensorType] = None,
         version: Optional[int] = None,
         **kwargs,
@@ -81,7 +84,9 @@ def __init__(
                 version = 2
         else:
             # old code
-            dims = _tensor_extra.infer_dim_tags(name=name, sparse_dim=sparse_dim, **kwargs)
+            dims = _tensor_extra.infer_dim_tags(
+                name=name, sparse_dim=sparse_dim, feature_dim_axis=feature_dim_axis, **kwargs
+            )
             if version is None:
                 version = 1
         if dtype is None:
@@ -96,6 +101,23 @@ def __init__(
         self.version = version
         self._extra = None  # type: Optional[_TensorExtra]
 
+        if feature_dim_axis is NotSpecified:
+            # version == 1 leaves it as NotSpecified,
+            #   which enables some dynamic default behavior, like taking the last axis.
+            # version >= 2 avoids any default behavior, and just sets it to None.
+            if version >= 2:
+                feature_dim_axis = None
+        elif feature_dim_axis is None:
+            pass
+        elif isinstance(feature_dim_axis, int):
+            assert not self.sparse_dim, "cannot have feature_dim_axis when sparse"
+            if feature_dim_axis < 0:
+                feature_dim_axis += self.batch_ndim
+            assert 0 <= feature_dim_axis < self.batch_ndim
+        else:
+            raise TypeError(f"unexpected feature_dim_axis type {type(feature_dim_axis)}")
+        self._feature_dim_axis = feature_dim_axis
+
         if kwargs:
             self._handle_extra_kwargs(**kwargs)
         if raw_tensor is not None:
@@ -123,5 +145,19 @@ def raw_tensor(self, value: Optional[RawTensorType]):
         self._raw_tensor = value
         self.sanity_check(assume_complete=False)
 
+    @property
+    def feature_dim(self) -> Optional[Dim]:
+        """
+        :return: self.dims[self.feature_dim_axis] or None
+        """
+        # first fast paths
+        if self._feature_dim_axis is None:
+            return None
+        if isinstance(self._feature_dim_axis, int):
+            return self._dims[self._feature_dim_axis]
+        if self.feature_dim_axis is None:
+            return None
+        return self._dims[self.feature_dim_axis]
+
 
 # dispatch table for frameworks for sanity_check


Test patch: