returnn 1.20240614.213333__tar.gz → 1.20240618.151147__tar.gz

{returnn-1.20240614.213333/returnn.egg-info → returnn-1.20240618.151147}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20240614.213333
+Version: 1.20240618.151147
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

returnn-1.20240618.151147/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240618.151147'
+long_version = '1.20240618.151147+git.84cb714'

{returnn-1.20240614.213333 → returnn-1.20240618.151147}/returnn/datasets/distrib_files.py

@@ -301,7 +301,7 @@ class DistributeFilesDataset(CachedDataset2):
         if "seq_ordering" not in dataset_dict and "seq_order_control_dataset" not in dataset_dict:
             raise ValueError(
                 f"{self}: sub dataset should have explicit seq_ordering "
-                "(or seq_order_control_dataset for MetaDataset), got: {dataset_dict}"
+                f"(or seq_order_control_dataset for MetaDataset), got: {dataset_dict}"
             )
         self._lazy_init_file_cache_proc()
         dataset_dict, exit_hook = self._file_cache.handle_cached_files_in_config(dataset_dict)

{returnn-1.20240614.213333 → returnn-1.20240618.151147}/returnn/frontend/__init__.py

@@ -21,6 +21,7 @@ from .state import *
 from .array_ import *
 from .attention import *
 from .backend import *
+from .build_from_dict import *
 from .cond import *
 from .const import *
 from .container import *
@@ -40,11 +41,13 @@ from .math_ import *
 from .matmul import *
 from .normalization import *
 from .parameter import *
+from .piecewise_linear import *
 from .rand import *
 from .rec import *
 from .reduce import *
 from .run_ctx import *
 from .signal import *
+from .stepwise_scheduler import *
 from .types import *
 
 # Modules not in the main namespace but in sub namespaces.

{returnn-1.20240614.213333 → returnn-1.20240618.151147}/returnn/frontend/_backend.py

@@ -135,14 +135,14 @@ class Backend(Generic[T]):
         raise NotImplementedError
 
     @staticmethod
-    def get_shape_raw(raw_tensor: T) -> Union[T, Tuple[Union[int, T]]]:
+    def get_shape_raw(raw_tensor: T) -> Union[T, Tuple[Union[int, T], ...]]:
         """
         :return: shape of raw tensor
         """
         raise NotImplementedError
 
     @staticmethod
-    def get_shape_tuple_raw(raw_tensor: T) -> Tuple[Union[int, T]]:
+    def get_shape_tuple_raw(raw_tensor: T) -> Tuple[Union[int, T], ...]:
         """
         :return: shape of raw tensor. assumes that ndim is known.
             In eager frameworks, all dims are int.
@@ -150,7 +150,7 @@ class Backend(Generic[T]):
         raise NotImplementedError
 
     @staticmethod
-    def get_known_shape_raw(raw_tensor: T) -> Tuple[Optional[int]]:
+    def get_known_shape_raw(raw_tensor: T) -> Tuple[Optional[int], ...]:
         """
         :return: shape of raw tensor, int for static known, None otherwise. assumes that ndim is known.
             This will not create any ops.
@@ -159,7 +159,7 @@ class Backend(Generic[T]):
         raise NotImplementedError
 
     @staticmethod
-    def set_known_shape_raw(raw_tensor: T, shape: Tuple[Optional[int]]) -> None:
+    def set_known_shape_raw(raw_tensor: T, shape: Tuple[Optional[int], ...]) -> None:
         """
         Sets the known shape of the raw tensor.
         This is only supported in graph-based frameworks,
@@ -506,6 +506,15 @@ class Backend(Generic[T]):
         """
         raise NotImplementedError
 
+    @staticmethod
+    def stack(sources: Sequence[Tensor], *, out_dim: Dim) -> Tensor:
+        """
+        :param sources:
+        :param out_dim:
+        :return: stacked tensor
+        """
+        raise NotImplementedError
+
     # Restrict the possible activation function names,
     # to not get unexpected behavior,
     # or unwanted incompatibilities.
@@ -953,6 +962,22 @@ class Backend(Generic[T]):
         """where"""
         raise NotImplementedError
 
+    @staticmethod
+    def search_sorted(
+        sorted_seq: Tensor, values: Tensor, *, axis: Dim, side: str = "left", out_dtype: str = "int32"
+    ) -> Tensor:
+        """
+        :param sorted_seq: [SharedDims...,axis], sequence of numbers, sorted low to high in the given axis.
+        :param values: [SharedDims...,OtherDims...], sequence of numbers to search for in ``sorted_seq``.
+        :param axis:
+        :param side: "left" or "right"
+        :param out_dtype:
+        :return: [SharedDims...,OtherDims...] -> axis, indices in axis in ``sorted_seq`` such that
+            sorted_seq[i-1] < value <= sorted_seq[i] if side=="left",
+            sorted_seq[i-1] <= value < sorted_seq[i] if side=="right".
+        """
+        raise NotImplementedError
+
     @staticmethod
     def clip_by_value(
         x: Tensor,
@@ -964,6 +989,25 @@ class Backend(Generic[T]):
         """clip by value"""
         raise NotImplementedError
 
+    @staticmethod
+    def lerp(
+        start: Tensor, end: Tensor, weight: Union[float, Tensor], *, allow_broadcast_all_sources: bool = False
+    ) -> Tensor:
+        """
+        Linear interpolation between start and end.
+        (Some backends might provide an optimized version of this.)
+
+        :param start:
+        :param end:
+        :param weight: scalar or tensor
+        :param allow_broadcast_all_sources:
+        :return: start + weight * (end - start)
+        """
+        # Default implementation.
+        if not allow_broadcast_all_sources:
+            return start + weight * (end - start)
+        return rf.combine_bc(start, "+", rf.combine_bc(weight, "*", rf.combine_bc(end, "-", start)))
+
     @staticmethod
     def matmul(a: Tensor[T], b: Tensor[T], *, reduce: Union[Dim, Sequence[Dim]], use_mask: bool = True) -> Tensor[T]:
         """

{returnn-1.20240614.213333 → returnn-1.20240618.151147}/returnn/frontend/_numpy_backend.py

@@ -52,14 +52,14 @@ class NumpyBackend(Backend[numpy.ndarray]):
         return raw_tensor.ndim
 
     @staticmethod
-    def get_shape_raw(raw_tensor: numpy.ndarray) -> Tuple[int]:
+    def get_shape_raw(raw_tensor: numpy.ndarray) -> Tuple[int, ...]:
         """
         :return: shape of raw tensor
         """
         return raw_tensor.shape
 
     @staticmethod
-    def get_shape_tuple_raw(raw_tensor: numpy.ndarray) -> Tuple[int]:
+    def get_shape_tuple_raw(raw_tensor: numpy.ndarray) -> Tuple[int, ...]:
         """
         :return: shape of raw tensor. assumes that ndim is known.
             In eager frameworks, all dims are int.
@@ -67,7 +67,7 @@ class NumpyBackend(Backend[numpy.ndarray]):
         return raw_tensor.shape
 
     @staticmethod
-    def get_known_shape_raw(raw_tensor: numpy.ndarray) -> Tuple[int]:
+    def get_known_shape_raw(raw_tensor: numpy.ndarray) -> Tuple[int, ...]:
         """
         :return: shape of raw tensor, int for static known, None otherwise. assumes that ndim is known.
             This will not create any ops.

{returnn-1.20240614.213333 → returnn-1.20240618.151147}/returnn/frontend/array_.py

@@ -29,6 +29,7 @@ __all__ = [
     "concat_features",
     "pad",
     "cum_concat_step",
+    "stack",
     "masked_select",
     "masked_scatter",
     "sequence_mask",
@@ -39,6 +40,7 @@ __all__ = [
     "shift_right",
     "reverse_sequence",
     "where",
+    "search_sorted",
     "sparse_to_dense",
     "one_hot",
 ]
@@ -352,6 +354,13 @@ def concat(
 ) -> Tuple[Tensor, Dim]:
     """
     Concatenates multiple sources in the specified dimension.
+
+    Also see :func:`stack`.
+
+    :param sources: list of (tensor, dim) pairs. dim is the axis to concatenate on.
+    :param allow_broadcast: if True, the sources can have different dims, and the result will be broadcasted.
+    :param out_dim: reuse existing dim for the resulting concatenated dim, if given
+    :return: concatenated tensor, out_dim
     """
     assert sources
     if not allow_broadcast:
@@ -490,6 +499,23 @@ def cum_concat_step(
     )
 
 
+def stack(sources: Sequence[Tensor], *, out_dim: Optional[Dim] = None) -> Tuple[Tensor, Dim]:
+    """
+    Stack the sources in a new dimension.
+    All sources must have the same shape.
+
+    :param sources:
+    :param out_dim: if given, use this as the new dim
+    :return: stacked tensor, out_dim
+    """
+    if not sources:
+        raise ValueError("no sources to stack")
+    if not out_dim:
+        out_dim = Dim(len(sources), name="stack")
+    # noinspection PyProtectedMember
+    return sources[0]._raw_backend.stack(sources, out_dim=out_dim), out_dim
+
+
 def masked_select(
     tensor: Tensor, *, mask: Tensor, dims: Sequence[Dim], out_dim: Optional[Dim] = None
 ) -> Tuple[Tensor, Dim]:
@@ -749,6 +775,23 @@ def where(
     return cond._raw_backend.where(cond, true_, false_, allow_broadcast_all_sources=allow_broadcast_all_sources)
 
 
+def search_sorted(
+    sorted_seq: Tensor, values: Tensor, *, axis: Dim, side: str = "left", out_dtype: str = "int32"
+) -> Tensor:
+    """
+    :param sorted_seq: [SharedDims...,axis], sequence of numbers, sorted low to high in the given axis.
+    :param values: [SharedDims...,OtherDims...], sequence of numbers to search for in ``sorted_seq``.
+    :param axis:
+    :param side: "left" or "right"
+    :param out_dtype:
+    :return: [SharedDims...,OtherDims...] -> axis, indices in axis in ``sorted_seq`` such that
+        sorted_seq[i-1] < value <= sorted_seq[i] if side=="left",
+        sorted_seq[i-1] <= value < sorted_seq[i] if side=="right".
+    """
+    # noinspection PyProtectedMember
+    return sorted_seq._raw_backend.search_sorted(sorted_seq, values, axis=axis, side=side, out_dtype=out_dtype)
+
+
 def sparse_to_dense(
     labels: Union[Tensor, rf.RawTensorTypes],
     *,

returnn-1.20240618.151147/returnn/frontend/build_from_dict.py

@@ -0,0 +1,60 @@
+"""
+Construct modules (or other objects) from dictionaries.
+"""
+
+from __future__ import annotations
+from typing import Union, Any, Type, Dict
+import importlib
+import returnn.frontend as rf
+
+
+__all__ = ["build_from_dict", "build_dict"]
+
+
+def build_from_dict(d: Dict[str, Any], *args, **kwargs) -> Union[rf.Module, Any]:
+    """
+    Build a module (or other object) from a dictionary.
+    `"class"` in the dict is required and specifies the class to be instantiated.
+    The other options are passed to the class constructor.
+
+    :param d: dictionary with the class name and other options
+    :param args: passed to the class constructor
+    :param kwargs: passed to the class constructor
+    :return: cls(*args, **d, **kwargs)
+    """
+    if "class" not in d:
+        raise ValueError(f"build_from_dict: Missing 'class' key in dict: {d}")
+    d = d.copy()
+    cls_name = d.pop("class")
+    cls = _get_cls(cls_name)
+    return cls(*args, **d, **kwargs)
+
+
+def build_dict(cls: Type, **kwargs) -> Dict[str, Any]:
+    """
+    Build a dictionary for :func:`build_from_dict`.
+    The class name is stored in the `"class"` key.
+
+    Note that this is intended to be used for serialization
+    and also to get a unique stable hashable representation
+    (e.g. for Sisyphus :func:`sis_hash_helper`)
+    which should not change if the class is renamed or moved
+    to keep the hash stable.
+    """
+    return {"class": _get_cls_name(cls), **kwargs}
+
+
+def _get_cls(cls_name: str) -> Type:
+    if "." not in cls_name:
+        raise ValueError(f"Expected '.' in class name: {cls_name}")
+    mod_name, cls_name = cls_name.rsplit(".", 1)
+    if mod_name == "rf":
+        return getattr(rf, cls_name)
+    mod = importlib.import_module(mod_name)
+    return getattr(mod, cls_name)
+
+
+def _get_cls_name(cls: Type) -> str:
+    if getattr(rf, cls.__name__, None) is cls:
+        return f"rf.{cls.__name__}"
+    return f"{cls.__module__}.{cls.__name__}"

{returnn-1.20240614.213333 → returnn-1.20240618.151147}/returnn/frontend/encoder/conformer.py

@@ -227,6 +227,10 @@ class ConformerEncoderLayer(rf.Module):
             conv_norm = rf.BatchNorm(out_dim, **conv_norm_opts)
         elif isinstance(conv_norm, type):
             conv_norm = conv_norm(out_dim, **(conv_norm_opts or {}))
+        elif isinstance(conv_norm, dict):
+            conv_norm = rf.build_from_dict(conv_norm, out_dim, **(conv_norm_opts or {}))
+        if not callable(conv_norm):
+            raise TypeError(f"{self}: unexpected conv_norm type {conv_norm!r}")
         self.conv_block = ConformerConvBlock(out_dim=out_dim, kernel_size=conv_kernel_size, norm=conv_norm)
         self.conv_layer_norm = rf.LayerNorm(out_dim)
 

{returnn-1.20240614.213333 → returnn-1.20240618.151147}/returnn/frontend/math_.py

@@ -66,6 +66,7 @@ __all__ = [
     "softmax",
     "log_softmax",
     "gating",
+    "lerp",
 ]
 
 
@@ -567,3 +568,20 @@ def gating(
 
     a, b = rf.split(x, axis=axis, out_dims=[out_dim, out_dim])
     return act_func(a) * gate_func(b), out_dim
+
+
+def lerp(
+    start: Tensor, end: Tensor, weight: Union[float, Tensor], *, allow_broadcast_all_sources: bool = False
+) -> Tensor:
+    """
+    Linear interpolation between start and end.
+    (Some backends might provide an optimized version of this.)
+
+    :param start:
+    :param end:
+    :param weight: scalar or tensor
+    :param allow_broadcast_all_sources:
+    :return: start + weight * (end - start)
+    """
+    # noinspection PyProtectedMember
+    return start._raw_backend.lerp(start, end, weight, allow_broadcast_all_sources=allow_broadcast_all_sources)

{returnn-1.20240614.213333 → returnn-1.20240618.151147}/returnn/frontend/normalization.py

@@ -8,7 +8,7 @@ from returnn.tensor import Tensor, Dim
 import returnn.frontend as rf
 
 
-__all__ = ["moments", "LayerNorm", "BatchNorm", "normalize", "Normalize"]
+__all__ = ["moments", "LayerNorm", "GroupNorm", "BatchNorm", "normalize", "Normalize"]
 
 
 def moments(
@@ -70,6 +70,30 @@ class LayerNorm(rf.Module):
         return norm_x * self.scale + self.bias
 
 
+class GroupNorm(rf.Module):
+    """
+    `Group normalization <https://arxiv.org/abs/1803.08494>`__.
+    """
+
+    def __init__(self, in_dim: Union[rf.Dim, Sequence[rf.Dim]], *, num_groups: Union[int, Dim], eps: float = 1e-6):
+        super().__init__()
+        self.in_dim = in_dim
+        self.num_groups = num_groups if isinstance(num_groups, Dim) else Dim(num_groups, name="groups")
+        self.in_group_dim = in_dim.ceildiv_left(num_groups)
+        self.eps = eps
+        self.scale = rf.Parameter([self.in_dim] if isinstance(self.in_dim, rf.Dim) else self.in_dim)
+        self.scale.initial = 1.0
+        self.bias = rf.Parameter(self.scale.dims)
+        self.bias.initial = 0.0
+
+    def __call__(self, x: Tensor) -> Tensor:
+        x = rf.split_dims(x, axis=self.in_dim, dims=[self.num_groups, self.in_group_dim])
+        mean, variance = rf.moments(x, axis=self.in_group_dim)
+        norm_x = (x - mean) * rf.rsqrt(variance + self.eps)
+        norm_x, _ = rf.merge_dims(norm_x, dims=[self.num_groups, self.in_group_dim], out_dim=self.in_dim)
+        return norm_x * self.scale + self.bias
+
+
 class BatchNorm(rf.Module):
     """
     Batch normalization. https://arxiv.org/abs/1502.03167

returnn-1.20240618.151147/returnn/frontend/piecewise_linear.py

@@ -0,0 +1,55 @@
+"""
+Piecewise linear function
+"""
+
+from __future__ import annotations
+from typing import Union, Dict
+import numpy as np
+from returnn.tensor import Tensor, Dim
+import returnn.frontend as rf
+
+
+__all__ = ["PiecewiseLinear"]
+
+
+class PiecewiseLinear(rf.Module):
+    """
+    Piecewise linear function.
+    """
+
+    def __init__(self, points: Dict[Union[int, float], Union[float, Tensor]]):
+        """
+        :param points: dict of key -> value pairs.
+        """
+        super().__init__()
+        if not points:
+            raise ValueError(f"{self}: points must not be empty")
+        self._points_sorted = sorted(points.items())
+        self.points_dim = Dim(len(self._points_sorted), name="pcw_schd_pieces")
+        # Note: Use rf.Parameter to work around deepcopy issue. https://github.com/rwth-i6/returnn/issues/1541
+        self._keys = rf.Parameter(
+            rf.convert_to_tensor(
+                np.array([k for k, _ in self._points_sorted], dtype=rf.get_default_float_dtype()),
+                dims=[self.points_dim],
+            ),
+            auxiliary=True,
+        )
+        self._values = rf.Parameter(
+            rf.stack([rf.convert_to_tensor(v) for _, v in self._points_sorted], out_dim=self.points_dim)[0],
+            auxiliary=True,
+        )
+
+    def __call__(self, x: Tensor) -> Tensor:
+        """
+        :param x: (x_dims...) -> value in keys
+        :return: y: (x_dims...,y_dims...) -> value in values
+        """
+        index = rf.search_sorted(self._keys, x, axis=self.points_dim)
+        index = rf.clip_by_value(index, 1, self.points_dim.dimension - 1)
+        x_start = rf.gather(self._keys, indices=index - 1)
+        x_end = rf.gather(self._keys, indices=index)
+        x_frac = (x - x_start) / (x_end - x_start)
+        x_frac = rf.clip_by_value(x_frac, 0.0, 1.0)
+        y_start = rf.gather(self._values, indices=index - 1)
+        y_end = rf.gather(self._values, indices=index)
+        return rf.lerp(y_start, y_end, x_frac)

{returnn-1.20240614.213333 → returnn-1.20240618.151147}/returnn/frontend/run_ctx.py

@@ -14,7 +14,7 @@ import returnn.frontend as rf
 from . import _backend
 
 
-__all__ = ["RunCtx", "Loss", "get_run_ctx", "init_train_step_run_ctx", "init_forward_step_run_ctx"]
+__all__ = ["RunCtx", "Loss", "get_run_ctx", "get_run_ctx_step", "init_train_step_run_ctx", "init_forward_step_run_ctx"]
 
 
 _run_ctx = None  # type: Optional[RunCtx]
@@ -134,6 +134,12 @@ class RunCtx:
         """
         return self._step
 
+    def get_step_tensor(self) -> Tensor:
+        """
+        :return: step as tensor
+        """
+        return rf.convert_to_tensor(self.step)
+
     @property
     def epoch(self) -> Union[int, Tensor]:
         """
@@ -141,6 +147,12 @@ class RunCtx:
         """
         return self._epoch
 
+    def get_epoch_tensor(self) -> Tensor:
+        """
+        :return: epoch as tensor
+        """
+        return rf.convert_to_tensor(self.epoch)
+
     def mark_as_loss(
         self,
         loss: Union[Tensor, Any],
@@ -350,6 +362,13 @@ class RunCtx:
         return loss
 
 
+def get_run_ctx_step() -> Tensor:
+    """
+    :return: shortcut for ``get_run_ctx().get_step_tensor()``
+    """
+    return get_run_ctx().get_step_tensor()
+
+
 @dataclass
 class Loss:
     """

returnn-1.20240618.151147/returnn/frontend/stepwise_scheduler.py

@@ -0,0 +1,31 @@
+"""
+Stepwise scheduler, e.g. for learning rate or other hyperparameters.
+
+All these modules will accept any args/kwargs but leave them unused,
+and instead uses :func:`get_run_ctx` to get the current train step from the current run context.
+"""
+
+from __future__ import annotations
+from returnn.tensor import Tensor
+import returnn.frontend as rf
+from .piecewise_linear import PiecewiseLinear
+
+
+__all__ = ["PiecewiseLinearStepwiseScheduler"]
+
+
+class PiecewiseLinearStepwiseScheduler(PiecewiseLinear):
+    """
+    Piecewise linear scheduler based on the current global train step.
+
+    Example::
+
+        scheduler = PiecewiseLinearStepwiseScheduler(
+            {0: 1.0, 10000: 0.1, 20000: 0.01}
+        )
+
+    This will start with 1.0, and then linearly decay to 0.1 at step 10000, and then to 0.01 at step 20000.
+    """
+
+    def __call__(self, *args, **kwargs) -> Tensor:
+        return super().__call__(rf.get_run_ctx_step())

{returnn-1.20240614.213333 → returnn-1.20240618.151147}/returnn/tensor/_dim_extra.py

@@ -4,7 +4,7 @@ or just rarely used attribs, such that we can save memory for the common case.
 """
 
 from __future__ import annotations
-from typing import TYPE_CHECKING, Optional, Union, Tuple, Sequence, Dict, List, Callable
+from typing import TYPE_CHECKING, Optional, Union, Any, Tuple, Sequence, Dict, List, Set, Callable
 import operator
 
 from returnn.util.basic import Entity
@@ -1286,15 +1286,15 @@ class _DimMixin:
         self: Dim,
         other: Dim,
         *,
-        ignore_feature_dim=False,
-        allow_same_feature_dim=False,
-        allow_same_spatial_dim=None,
-        treat_feature_as_spatial=False,
-        broadcast_matches=False,
-        unknown_spatial_matches=False,
-        undefined_matches=False,
-        derived_matches=False,
-        allow_old_behavior=False,
+        ignore_feature_dim: bool = False,
+        allow_same_feature_dim: bool = False,
+        allow_same_spatial_dim: Optional[bool] = None,
+        treat_feature_as_spatial: bool = False,
+        broadcast_matches: bool = False,
+        unknown_spatial_matches: bool = False,
+        undefined_matches: bool = False,
+        derived_matches: bool = False,
+        allow_old_behavior: bool = False,
     ) -> bool:
         """
         Compares self to other for equality.
@@ -1307,16 +1307,16 @@ class _DimMixin:
         and might potentially change in the future.
           https://github.com/rwth-i6/returnn/issues/634
 
-        :param Dim other:
-        :param bool ignore_feature_dim:
-        :param bool allow_same_feature_dim:
-        :param bool|None allow_same_spatial_dim:
-        :param bool treat_feature_as_spatial:
-        :param bool broadcast_matches:
-        :param bool unknown_spatial_matches:
-        :param bool undefined_matches:
-        :param bool derived_matches:
-        :param bool allow_old_behavior: useful e.g. for find_matching_dim_map
+        :param other:
+        :param ignore_feature_dim:
+        :param allow_same_feature_dim:
+        :param allow_same_spatial_dim:
+        :param treat_feature_as_spatial:
+        :param broadcast_matches:
+        :param unknown_spatial_matches:
+        :param undefined_matches:
+        :param derived_matches:
+        :param allow_old_behavior: useful e.g. for find_matching_dim_map
         """
         if self is other:  # first some fast path check
             return True
@@ -1820,12 +1820,13 @@ class _DimMixin:
         self._make_extra().copy_same_as = other
 
     @classmethod
-    def get_existing_tag_from_collection(cls, other, tags, is_equal_opts=None):
+    def get_existing_tag_from_collection(
+        cls, other: Dim, tags: Union[Sequence[Dim], Set[Dim]], is_equal_opts: Optional[Dict[str, Any]] = None
+    ) -> Optional[Dim]:
         """
-        :param Dim other:
-        :param list[Dim]|tuple[Dim]|set[Dim] tags:
-        :param dict[str]|None is_equal_opts: passed to Dim.is_equal
-        :rtype: Dim|None
+        :param other:
+        :param tags:
+        :param is_equal_opts: passed to Dim.is_equal
         """
         if is_equal_opts is None:
             is_equal_opts = {}
@@ -1842,13 +1843,17 @@ class _DimMixin:
         return None
 
     @classmethod
-    def get_all_dimension_tags(cls, data_list, is_equal_opts=None, unique_separate_axes=True):
-        """
-        :param list[_t.Tensor] data_list:
-        :param dict[str]|None is_equal_opts: passed to Dim.is_equal
-        :param bool unique_separate_axes: e.g. data_list=[Data with shape (B,5,5,10)] results in 4 dim tags, not 3.
+    def get_all_dimension_tags(
+        cls,
+        data_list: List[_t.Tensor],
+        is_equal_opts: Optional[Dict[str, Any]] = None,
+        unique_separate_axes: bool = True,
+    ) -> Tuple[List[Dim], util.DictRefKeys[_t.Tensor, List[Dim]]]:
+        """
+        :param data_list:
+        :param is_equal_opts: passed to Dim.is_equal
+        :param unique_separate_axes: e.g. data_list=[Data with shape (B,5,5,10)] results in 4 dim tags, not 3.
         :return: list of dimension tags, dict for data -> list of dimension tags (for each axis)
-        :rtype: (list[Dim], util.DictRefKeys[_t.Tensor, list[Dim]])
         """
         tags = []
         data_axes_dict = util.DictRefKeys()  # type: util.DictRefKeys[_t.Tensor, List[Dim]]
@@ -2323,11 +2328,11 @@ class Op:
     Op on :class:`Dim` which results in a derived :class:`Dim`.
     """
 
-    def __init__(self, kind, inputs, attribs=None):
+    def __init__(self, kind: str, inputs: List[Dim], attribs: Optional[Dict[str, Any]] = None):
         """
-        :param str kind: "add", "sub", "mul", "ceildiv"
-        :param list[Dim] inputs:
-        :param dict[str]|None attribs:
+        :param kind: "add", "sub", "mul", "ceildiv"
+        :param inputs:
+        :param attribs:
         """
         self.kind = kind
         self.inputs = inputs

{returnn-1.20240614.213333 → returnn-1.20240618.151147}/returnn/tf/frontend_layers/_backend.py

@@ -151,12 +151,12 @@ class ReturnnLayersBackend(Backend[Layer]):
         raise NotImplementedError
 
     @staticmethod
-    def get_shape_tuple_raw(raw_tensor: Layer) -> Tuple[Union[int, Layer]]:
+    def get_shape_tuple_raw(raw_tensor: Layer) -> Tuple[Union[int, Layer], ...]:
         """shape"""
         raise NotImplementedError
 
     @staticmethod
-    def get_known_shape_raw(raw_tensor: Layer) -> Tuple[Optional[int]]:
+    def get_known_shape_raw(raw_tensor: Layer) -> Tuple[Optional[int], ...]:
         """known shape"""
         return raw_tensor.tensor.batch_shape