returnn 1.20250220.174943__py3-none-any.whl → 1.20250221.114352__py3-none-any.whl

returnn/PKG-INFO

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

returnn/_setup_info_generated.py

@@ -1,2 +1,2 @@
-version = '1.20250220.174943'
-long_version = '1.20250220.174943+git.1766d1c'
+version = '1.20250221.114352'
+long_version = '1.20250221.114352+git.650b638'

returnn/frontend/_backend.py

@@ -784,18 +784,11 @@ class Backend(Generic[T]):
         dims: Sequence[Dim],
         dtype: str,
         sparse_dim: Optional[Dim] = None,
+        feature_dim: Optional[Dim] = None,
         device: Optional[str] = None,
         name: Optional[str] = None,
     ) -> Tensor[T]:
-        """
-        :param value: tensor, or scalar raw tensor or some other scalar value
-        :param dims:
-        :param dtype:
-        :param sparse_dim:
-        :param device:
-        :param name:
-        :return: tensor
-        """
+        """convert (raw/any) tensor to tensor"""
         raise NotImplementedError
 
     @staticmethod
@@ -956,6 +949,11 @@ class Backend(Generic[T]):
         """where"""
         raise NotImplementedError
 
+    @staticmethod
+    def sort(source: Tensor, *, axis: Dim, descending: bool, stable: bool) -> Tuple[Tensor, Tensor, Dim]:
+        """sort. return values and indices"""
+        raise NotImplementedError
+
     @staticmethod
     def search_sorted(
         sorted_seq: Tensor, values: Tensor, *, axis: Dim, side: str = "left", out_dtype: str = "int32"

returnn/frontend/_native/module.cpp

@@ -312,8 +312,9 @@ bool PyModuleState::_cachedOpInitTorch() {
     AddOp(TOp_FloorDiv, "floor_divide");
     AddOp(TOp_Mod, "remainder");
     AddOp(TOp_Pow, "pow");
-    AddOp(TOp_Maximum, "maximum");
-    AddOp(TOp_Minimum, "minimum");
+    // Use clamp_min/clamp_max instead of maximum/minimum because the former allow number arguments.
+    AddOp(TOp_Maximum, "clamp_min");
+    AddOp(TOp_Minimum, "clamp_max");
     AddOpAlt(TOp_SquaredDifference, "squared_difference");
     AddOp(TOp_And, "logical_and");
     AddOp(TOp_Or, "logical_or");

returnn/frontend/_native/tensor_ops.cpp

@@ -1368,6 +1368,14 @@ static PyObject* compareOrCombineViaCached(
         case TOp_FloorDiv:
         case TOp_Mod:
         case TOp_Pow:
+        case TOp_Maximum:
+        case TOp_Minimum:
+        case TOp_Eq:
+        case TOp_Ne:
+        case TOp_Lt:
+        case TOp_Le:
+        case TOp_Gt:
+        case TOp_Ge:
             needConvertToTensor = false;
         default:
             break;

returnn/frontend/_numpy_backend.py

@@ -83,6 +83,7 @@ class NumpyBackend(Backend[numpy.ndarray]):
         dims: Sequence[Dim],
         dtype: str,
         sparse_dim: Optional[Dim] = None,
+        feature_dim: Optional[Dim] = None,
         device: Optional[str] = None,
         name: Optional[str] = None,
     ) -> Tensor[numpy.ndarray]:
@@ -95,7 +96,7 @@ class NumpyBackend(Backend[numpy.ndarray]):
             name = name or "const"
             value = numpy.array(value, dtype=NumpyBackend.as_dtype_raw(dtype))
         assert isinstance(value, numpy.ndarray)
-        return Tensor(name, dims=dims, dtype=dtype, sparse_dim=sparse_dim, raw_tensor=value)
+        return Tensor(name, dims=dims, dtype=dtype, sparse_dim=sparse_dim, feature_dim=feature_dim, raw_tensor=value)
 
     @staticmethod
     def expand_dims_raw(raw_tensor: numpy.ndarray, axis: int) -> numpy.ndarray:

returnn/frontend/array_.py

@@ -45,9 +45,12 @@ __all__ = [
     "shift_left",
     "reverse_sequence",
     "where",
+    "sort",
     "search_sorted",
     "sparse_to_dense",
     "one_hot",
+    "top_k_mask",
+    "top_p_mask",
 ]
 
 
@@ -57,6 +60,7 @@ def convert_to_tensor(
     dims: Sequence[Dim] = None,
     dtype: Optional[str] = None,
     sparse_dim: Optional[Dim] = None,
+    feature_dim: Optional[Dim] = None,
     shape: Sequence[Dim] = None,
     device: Optional[str] = None,
     keep_scalar_on_cpu: bool = False,
@@ -68,6 +72,7 @@ def convert_to_tensor(
     :param dims:
     :param dtype:
     :param sparse_dim:
+    :param feature_dim:
     :param shape: alias for dims, for some older code
     :param name:
     :param device:
@@ -121,7 +126,7 @@ def convert_to_tensor(
         if dtype is None:
             dtype = value_backend.get_dtype_name_raw(value)
     return _backend.convert_to_tensor(
-        value=value, dims=dims, dtype=dtype, sparse_dim=sparse_dim, device=device, name=name
+        value=value, dims=dims, dtype=dtype, sparse_dim=sparse_dim, feature_dim=feature_dim, device=device, name=name
     )
 
 
@@ -996,6 +1001,27 @@ def where(
     return cond._raw_backend.where(cond, true_, false_, allow_broadcast_all_sources=allow_broadcast_all_sources)
 
 
+def sort(source: Tensor, *, axis: Dim, descending: bool = False, stable: bool = True) -> Tuple[Tensor, Tensor, Dim]:
+    """
+    Sorts the source tensor along the given axis.
+
+    See also :func:`top_k`.
+    :func:`top_k` with ``k=axis.get_size_tensor()`` is equivalent to this function.
+
+    :param source: {other_dims..., axis}
+    :param axis: The axis to sort along.
+    :param descending: If True, sort in descending order, otherwise in ascending order.
+    :param stable: If True, use a stable sorting algorithm (not reordering equal elements).
+        Note that many frameworks (Torch, TensorFlow) have ``stable=False`` by default.
+        ``stable=False`` can be faster.
+    :return: sorted tensor, indices tensor, out_dim. both tensors have the shape {other_dims..., out_dim},
+        i.e. ``axis`` replaced by ``out_dim``.
+        indices tensor has sparse_dim set to ``axis``.
+    """
+    # noinspection PyProtectedMember
+    return source._raw_backend.sort(source, axis=axis, descending=descending, stable=stable)
+
+
 def search_sorted(
     sorted_seq: Tensor, values: Tensor, *, axis: Dim, side: str = "left", out_dtype: str = "int32"
 ) -> Tensor:
@@ -1044,3 +1070,49 @@ def one_hot(source: Tensor) -> Tensor:
     and much more efficiently than they would be with dense tensors.
     """
     return sparse_to_dense(source, label_value=1.0, other_value=0.0)
+
+
+def top_k_mask(values: Tensor, *, axis: Dim, k: Union[int, Tensor]) -> Tensor:
+    """
+    Top-k filtering.
+
+    :param values: {other_dims..., axis}
+    :param axis:
+    :param k: the number of top values to keep
+    :return: mask {other_dims..., axis} of the top-k values
+    """
+    _, indices, k_dim = rf.top_k(values, axis=axis, k=k)
+    mask = rf.scatter(rf.full(dims=indices.dims, fill_value=True), indices=indices, indices_dim=k_dim, fill_value=False)
+    return mask
+
+
+def top_p_mask(
+    probs: Tensor,
+    *,
+    axis: Dim,
+    p: Union[float, Tensor],
+    one_more: bool = True,
+) -> Tensor:
+    """
+    Top-p filtering, e.g. as used in Nucleus sampling (https://arxiv.org/abs/1904.09751).
+
+    :param probs: {probs_dims..., axis}
+    :param axis:
+    :param p: the probability mass to keep
+    :param one_more: if True (default), keep also the first token above the threshold.
+        (It's enabled by default to follow the behavior of the original implementation.)
+    :return: mask {probs_dims..., axis} of the top-p tokens.
+        ``sum(probs[mask]) <= p``, or slightly more if ``one_more`` is True.
+    """
+    assert 0.0 <= p <= 1.0
+    if isinstance(p, Tensor):
+        assert axis not in p.dims
+    # https://github.com/ari-holtzman/degen/blob/master/gen.py
+    sorted_probs, sorted_indices, sorted_dim = rf.sort(probs, axis=axis, descending=True)
+    cum_probs = rf.cumsum(sorted_probs, spatial_dim=sorted_dim)
+    mask = cum_probs <= p  # {probs_dims..., sorted_dim}
+    if one_more:
+        # keep also the first token above the threshold
+        mask = rf.shift_right(mask, axis=sorted_dim, pad_value=True)
+    mask = rf.scatter(mask, indices=sorted_indices, indices_dim=sorted_dim)
+    return mask

returnn/frontend/dims.py

@@ -14,6 +14,7 @@ __all__ = [
     "range_over_dim",
     "range_over_dim_strided",
     "range_over_merged_dims",
+    "linspace_over_dim",
     "replace_dim",
     "replace_dim_v2",
     "set_sparse_dim",
@@ -81,6 +82,36 @@ def range_over_merged_dims(
     return indices
 
 
+def linspace_over_dim(
+    dim: Dim,
+    start: Union[float, Tensor] = 0.0,
+    end: Union[float, Tensor] = 1.0,
+    *,
+    dtype: Optional[str] = None,
+    device: Optional[str] = None,
+) -> Tensor:
+    """
+    Linearly spaced values over a dim.
+
+    :param dim: dim to range over
+    :param start: start value
+    :param end: end value
+    :param dtype: dtype of the output tensor
+    :param device: device of the output tensor
+    :return: tensor with shape [dim] containing linearly spaced values between start and end
+    """
+    if dtype is None:
+        dtype = rf.get_default_float_dtype()
+    indices = rf.range_over_dim(dim, dtype=dtype, device=device)
+    linspace = indices / rf.cast(rf.maximum(dim.get_size_tensor(device=indices.device), 1), dtype=indices.dtype)
+    space_len = end - start
+    if not isinstance(space_len, (int, float)) or space_len != 1:
+        linspace *= space_len
+    if not isinstance(start, (int, float)) or start != 0:
+        linspace += start
+    return linspace
+
+
 def replace_dim(source: Tensor, *, in_dim: Dim, out_dim: Optional[Dim] = None) -> Tuple[Tensor, Dim]:
     """
     Also see: :func:`replace_dim_v2`, :func:`rf.merge_dims`, :func:`rf.split_dims`.

returnn/frontend/rand.py

@@ -64,6 +64,7 @@ __all__ = [
     "random_uniform",
     "random_normal",
     "random_truncated_normal",
+    "random_choice_without_replacement",
 ]
 
 
@@ -349,3 +350,32 @@ def random_truncated_normal(
         static=static,
         out=out,
     )
+
+
+def random_choice_without_replacement(
+    *,
+    log_probs: Tensor,
+    axis: Union[Dim, Sequence[Dim]],
+    num_samples_dim: Dim,
+    noise_scale: Union[float, Tensor] = 1.0,
+) -> Union[Tensor, Sequence[Tensor]]:
+    """
+    Randomly sample without replacement.
+
+    :param log_probs: {log_probs_dims..., axis}
+    :param axis: same as in :func:`top_k`
+    :param num_samples_dim: how many samples to draw
+    :param noise_scale: scale the noise. with scale=0, you get :func:`top_k`.
+    :return: random indices shape {log_probs_dims..., num_samples_dim} -> axis.
+        if axis was a sequence, will return a sequence of tensors.
+    """
+    # https://github.com/tensorflow/tensorflow/issues/9260
+    # https://timvieira.github.io/blog/post/2014/08/01/gumbel-max-trick-and-weighted-reservoir-sampling/
+    scores_random_sample = -rf.log(
+        -rf.log(random_uniform(log_probs.dims, dtype=log_probs.dtype, device=log_probs.device))
+    )
+    if not isinstance(noise_scale, (int, float)) or noise_scale != 1.0:
+        scores_random_sample *= noise_scale
+    scores = log_probs + scores_random_sample
+    _, indices, _ = rf.top_k(scores, k_dim=num_samples_dim, axis=axis)
+    return indices

returnn/tf/frontend_layers/_backend.py

@@ -559,6 +559,7 @@ class ReturnnLayersBackend(Backend[Layer]):
         dims: Sequence[Dim],
         dtype: str,
         sparse_dim: Optional[Dim] = None,
+        feature_dim: Optional[Dim] = None,
         device: Optional[str] = None,
         name: Optional[str] = None,
     ) -> Tensor[Layer]:
@@ -568,6 +569,8 @@ class ReturnnLayersBackend(Backend[Layer]):
         kwargs = {}
         if sparse_dim:
             kwargs["sparse_dim"] = sparse_dim
+        if feature_dim:
+            kwargs["feature_dim"] = feature_dim
         dim_deps = _dims.get_dim_deps(dims)
         if dim_deps:
             kwargs["shape_deps"] = dim_deps

returnn/tf/frontend_low_level/_backend.py

@@ -411,24 +411,19 @@ class TFBackend(Backend[tf.Tensor]):
         dims: Sequence[Dim],
         dtype: str,
         sparse_dim: Optional[Dim] = None,
+        feature_dim: Optional[Dim] = None,
         device: Optional[str] = None,
         name: Optional[str] = None,
     ) -> _TT:
-        """
-        :param value:
-        :param dims:
-        :param dtype:
-        :param sparse_dim:
-        :param device:
-        :param name:
-        :return: tensor
-        """
+        """convert to tensor"""
         if isinstance(value, Tensor):
             return value
         with tf.control_dependencies(None):
             value = tf.convert_to_tensor(value, dtype=dtype)
         assert isinstance(value, tf.Tensor)
-        return Tensor(name or "const", raw_tensor=value, dims=dims, dtype=dtype, sparse_dim=sparse_dim)
+        return Tensor(
+            name or "const", raw_tensor=value, dims=dims, dtype=dtype, sparse_dim=sparse_dim, feature_dim=feature_dim
+        )
 
     @staticmethod
     def range_over_dim(dim: Dim, *, dtype: Optional[str] = None, device: Optional[str] = None) -> _TT:

returnn/torch/distributed.py

@@ -3,10 +3,11 @@ torch.distributed utils
 """
 
 from __future__ import annotations
-from typing import Optional, Any, Dict
+import ast
+import logging
 import os
 import socket
-import logging
+from typing import Optional, Any, Dict
 
 import torch
 from torch.nn.parallel import DistributedDataParallel
@@ -23,19 +24,31 @@ class DistributedContext:
     """
 
     def __init__(self, options: Dict[str, Any]):
-        import torch.distributed as dist
-
         self._opts = CollectionReadCheckCovered(options)
 
-        # when no backend is specified, both gloo and nccl backends will be created
-        # the gloo backend will be used for collectives with CPU tensors and
-        # the nccl backend will be used for collectives with CUDA tensors
-        dist.init_process_group(backend=self._opts.get("backend", None))
+        # Subprocesses have issues initializing torch.distributed process groups.
+        #
+        # We therefore pass rank/size information of the process group via an env
+        # variable that is automatically inherited in any created subprocess.
+        env_var_name = "_RETURNN_TORCH_DISTRIBUTED_INIT_INFO"
+        prev_init_info = os.environ.get(env_var_name)
+        if prev_init_info:
+            self.prev_init_info = ast.literal_eval(prev_init_info)
+            self._rank = self.prev_init_info["rank"]
+            self._size = self.prev_init_info["size"]
+        else:
+            import torch.distributed as dist
+
+            # when no backend is specified, both gloo and nccl backends will be created
+            # the gloo backend will be used for collectives with CPU tensors and
+            # the nccl backend will be used for collectives with CUDA tensors
+            dist.init_process_group(backend=self._opts.get("backend", None))
+            self._rank = dist.get_rank()
+            self._size = dist.get_world_size()
+            os.environ[env_var_name] = repr({"rank": self._rank, "size": self._size})
 
         self._local_rank = int(os.environ["LOCAL_RANK"])
         self._local_size = int(os.environ["LOCAL_WORLD_SIZE"])
-        self._rank = dist.get_rank()
-        self._size = dist.get_world_size()
 
         _logger.info(
             "Torch distributed initialized. Hostname %s, pid %i, rank %i / size %i, local rank %s / local size %s."
@@ -123,9 +136,9 @@ _is_set_up = False
 _ctx = None  # type: Optional[DistributedContext]
 
 
-def get_ctx(config=None) -> Optional[DistributedContext]:
+def get_ctx(config: Optional[Config] = None) -> Optional[DistributedContext]:
     """
-    :param Config|None config:
+    :param config:
     :returns: the global context if Torch distributed is enabled, or None otherwise.
       If we did not setup the context yet, it will automatically create it.
     """

returnn/torch/frontend/_backend.py

@@ -895,18 +895,11 @@ class TorchBackend(Backend[torch.Tensor]):
         dims: Sequence[Dim],
         dtype: str,
         sparse_dim: Optional[Dim] = None,
+        feature_dim: Optional[Dim] = None,
         device: Optional[str] = None,
         name: Optional[str] = None,
     ) -> Tensor[torch.Tensor]:
-        """
-        :param value:
-        :param dims:
-        :param dtype:
-        :param sparse_dim:
-        :param device:
-        :param name:
-        :return: tensor
-        """
+        """convert to tensor"""
         if isinstance(value, Tensor):
             return value
         if isinstance(value, torch.Tensor):
@@ -926,7 +919,7 @@ class TorchBackend(Backend[torch.Tensor]):
                     device=device or rf.get_default_device(),
                 )
         assert isinstance(value, torch.Tensor)
-        return Tensor(name, dims=dims, dtype=dtype, sparse_dim=sparse_dim, raw_tensor=value)
+        return Tensor(name, dims=dims, dtype=dtype, sparse_dim=sparse_dim, feature_dim=feature_dim, raw_tensor=value)
 
     @staticmethod
     def full(
@@ -1223,6 +1216,21 @@ class TorchBackend(Backend[torch.Tensor]):
         out.raw_tensor = torch.where(cond_bc_raw, true_bc_raw, false_bc_raw)
         return out
 
+    @staticmethod
+    def sort(source: Tensor, *, axis: Dim, descending: bool, stable: bool) -> Tuple[Tensor, Tensor, Dim]:
+        """sort. return values and indices"""
+        axis_int = source.get_axis_from_description(axis, allow_int=False)
+        # Move to last axis. Should be more efficient.
+        source = source.copy_move_axis(axis_int, -1)
+        axis_int = source.batch_ndim - 1
+        values_raw, indices_raw = torch.sort(source.raw_tensor, dim=axis_int, descending=descending, stable=stable)
+        out_dims = list(source.dims)
+        out_dim = axis.copy(same_as_self=False, description=f"{axis.description}:sorted")
+        out_dims[axis_int] = out_dim
+        values = rf.convert_to_tensor(values_raw, dims=out_dims, feature_dim={axis: out_dim}.get(source.feature_dim))
+        indices = rf.convert_to_tensor(indices_raw, dims=out_dims, sparse_dim=axis)
+        return values, indices, out_dim
+
     @staticmethod
     def search_sorted(
         sorted_seq: Tensor, values: Tensor, *, axis: Dim, side: str = "left", out_dtype: str = "int32"
@@ -1566,6 +1574,9 @@ class TorchBackend(Backend[torch.Tensor]):
             return values, indices_out, k_dim
         assert isinstance(axis, Dim)
         axis_int = source.get_axis_from_description(axis, allow_int=False)
+        # Move to last axis. Should be more efficient.
+        source = source.copy_move_axis(axis_int, -1)
+        axis_int = source.batch_ndim - 1
         values_raw, indices_raw = torch.topk(
             source.raw_tensor, k=k_dim.get_dim_value(), dim=axis_int, largest=True, sorted=sorted
         )

returnn/util/basic.py

@@ -3773,9 +3773,9 @@ def should_write_to_disk(config):
     if config.typed_value("torch_distributed") is not None:
         assert BackendEngine.is_torch_selected(), "torch_distributed assumes PyTorch"
 
-        import torch.distributed
+        import returnn.torch.distributed as torch_distributed
 
-        if torch.distributed.get_rank() != 0:
+        if torch_distributed.get_ctx(config).rank() != 0:
             return False
     elif config.is_true("use_horovod"):
         assert BackendEngine.is_tensorflow_selected(), "use_horovod currently assumes TensorFlow"

{returnn-1.20250220.174943.dist-info → returnn-1.20250221.114352.dist-info}/METADATA

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

{returnn-1.20250220.174943.dist-info → returnn-1.20250221.114352.dist-info}/RECORD

@@ -1,9 +1,9 @@
-returnn/PKG-INFO,sha256=SqIvwRtUoU4A4DTo3MA0lH5LDudxkcTNX0sV2dn-cNM,5215
+returnn/PKG-INFO,sha256=VL6JCmrRQYGi4OdodEygiq3hJyM1C-73zvMTSusYpOs,5215
 returnn/__init__.py,sha256=biBtRsM0WZ406vShaeH-9WFoqJ8XwTbn6g0EeFJ7l8E,1012
 returnn/__main__.py,sha256=qBFbuB1yN3adgVM5pXt2-Yq9vorjRNchNPL8kDKx44M,31752
 returnn/__old_mod_loader__.py,sha256=nvsNY-xELdS_IPNkv66Q9Rmvg4dbGW0-EBRDcCmctos,7654
 returnn/__setup__.py,sha256=22kQn2fh11iPM0hLb2Fy5sLmoU1JGvmDxXRYuRgQkwU,4659
-returnn/_setup_info_generated.py,sha256=6XGlRiImjQzRDEvGJh0r14QAZsL7lAgNWuILyYOtoyk,77
+returnn/_setup_info_generated.py,sha256=LlW75YDQH_DvPCMrSeF0bY52JGq9l4tJNA5mGTT5MQA,77
 returnn/config.py,sha256=3tmKhB6FnQZaNdtcYsiB61JnEY--iZ2qmJ4yq0b6tE0,29140
 returnn/forward_iface.py,sha256=A_OJiaXsX4MlXQRzST86ylyxSUZbC402PQL1REcqHjM,911
 returnn/learning_rate_control.py,sha256=ZvWryAn_tv9DhV8sh1LV3eE34Yltl3On3mYZAG4hR9s,34684
@@ -75,12 +75,12 @@ returnn/extern/graph_editor/subgraph.py,sha256=R3uIFqWgiL7L5S4YATm9o9a3wfEa_mSb4
 returnn/extern/graph_editor/transform.py,sha256=d9fEgu0JC342q0g9niVxRWMKzkQQA9mrrajBGcU1o_s,29349
 returnn/extern/graph_editor/util.py,sha256=QMrQeQZ7lJwsrNQub9tof0h3quEaoHiGJaZmogQ7jXE,18707
 returnn/frontend/__init__.py,sha256=2aS7nbxXniIrBp2DODl0xN0f3IJ_dX4Bi9ZlR7W5_DE,1472
-returnn/frontend/_backend.py,sha256=datTTGM6xhvpiSPphDZgRjruU-SxrowG15_ZnTkEKy0,50108
+returnn/frontend/_backend.py,sha256=W3J3ZSOxonX6wk-wY2dX_aokXHpm1VQ1V0qSjllQxUM,50165
 returnn/frontend/_cache.py,sha256=JAhi7L-raQ3A-NC3JUYDtdRTwT3BGJJGGZxrZ8MfEWQ,8403
-returnn/frontend/_numpy_backend.py,sha256=_akZZkZplzR8XNVIrenvR99s-CapWDYL9WSD7VXwGYI,7790
+returnn/frontend/_numpy_backend.py,sha256=2oCtG0YCWL_89v4cD_jDj8em1O_Fp-_YWl5EblGi_yo,7858
 returnn/frontend/_random_journal.py,sha256=_ktP_mjgx8vtQQGX_DofdhewJj0aPiczefTWeemPkmo,5457
 returnn/frontend/_utils.py,sha256=4A3MSRM0i86J77550uR_AjcBEPu6nymLUZ9Xd1V3Fkc,12073
-returnn/frontend/array_.py,sha256=jzYiWUhiQ7Ldl2ja81NlZ9SLh_zlPGTJLSmAE77bYrk,41236
+returnn/frontend/array_.py,sha256=UHTQmb_cFsjVStAELcCqMkCbQNQiBiwN4gQZu6CloIA,44126
 returnn/frontend/attention.py,sha256=GKt-Xqnz8sIyXVrE0i4VCS7J2Wu7dmoH_BA0Cu8CrXQ,45769
 returnn/frontend/backend.py,sha256=iQ9w4xl8Ea7bgpb0VUaCKq50rV5Bl2E5J8Rhd-oqD_c,883
 returnn/frontend/build_from_dict.py,sha256=rfWa2rjjhIR_kIQED_nMrygrQBunS6unegzWTLVbC98,3017
@@ -90,7 +90,7 @@ returnn/frontend/container.py,sha256=wF3OlQN7WlOVmmdapUth_Unha3DVf6h1B7okBJAuJDA
 returnn/frontend/control_flow_ctx.py,sha256=v17CsNwRnZYe8GdMtGJt2ftibfxMCGK1i0l-GX5ILu0,699
 returnn/frontend/conv.py,sha256=51LZovcRzITDLXvPcJs_MFsGEY_MFvO_MFF9D-jZstA,22481
 returnn/frontend/device.py,sha256=K7Y1qoQcO4GIHgLkPLQWK-GVT8gKL8GwyQrmPo8LgBE,1438
-returnn/frontend/dims.py,sha256=FK2BUVlftvyA8DRtujylUIES_Jw8QH3aB79wEmPchj4,9841
+returnn/frontend/dims.py,sha256=hKA7IQRB0DbohN1ngNw31W44BsyjdHCtYAccxOcumzQ,10872
 returnn/frontend/dropout.py,sha256=rsx3p5b0NblBfXXSQZTQFJ8jUUS3fj4Qzc39iffBMCA,5006
 returnn/frontend/dtype.py,sha256=Ooc5BrcNrTp6XShuFEV9g5V6-niuy4ImP_Lt_Qgq3jE,1886
 returnn/frontend/gradient.py,sha256=dOUvLqN-vxsvjKQfpfIvEYlx4TlpHkOk-p9hsB680iA,3376
@@ -110,7 +110,7 @@ returnn/frontend/parameter.py,sha256=w6SN-uv87OyeWBt90_3UBbK0h6sftSOCxkqXPg76caY
 returnn/frontend/parametrizations.py,sha256=hVbOlgm1pQAmZnAnNxq8Tk23rykr_iy3-6R1H6CwlMA,2798
 returnn/frontend/parametrize.py,sha256=VhgTEP7ehON950Q4bkCy8rvg9641moEKAXn0XzomK6E,7216
 returnn/frontend/piecewise_linear.py,sha256=TdL6wzop8P1dcIZwkEbJFvSUZSI1cbhS3XKzlWQkEVI,1964
-returnn/frontend/rand.py,sha256=uYnw3OvzMVZlNq-aHtX7Lvmb84UO7IVTcrSAlX3TbGg,11731
+returnn/frontend/rand.py,sha256=Levgf5VtOOBKDSgz0869Jf3VW4BWxYZuRXsa_fOxNI4,12969
 returnn/frontend/rec.py,sha256=4m20LvsPJ75pRYykVrup6Csj_D7duG-dW28SaJh-sq8,7863
 returnn/frontend/reduce.py,sha256=-Zt-OH6Zbtb9uR6YEzurCyrowH-anIXvuga6Pla2V70,10220
 returnn/frontend/run_ctx.py,sha256=ItcZwuFItkZjYWrg715L1Za2Xg7__MQCrRCAwBeTUxA,21411
@@ -122,10 +122,10 @@ returnn/frontend/types.py,sha256=gpevnXZSlF_BgA76duIkkzN-ed_MflhSlOnHj1xJnAs,111
 returnn/frontend/_native/__init__.py,sha256=fVjazAujt0rdICXZL-GgW1sjFeL1HB4NPuy2m5rmMsc,6480
 returnn/frontend/_native/backend.cpp,sha256=MeHczHypwj_ncntOxRqanK8SqGyV9Eq1X0cpMWb_WII,4768
 returnn/frontend/_native/backend.hpp,sha256=Wq80dcEzXfRNxGOXFnIgHllkiv1rDi3KpHK-xxJsSDI,791
-returnn/frontend/_native/module.cpp,sha256=G-U9OAKlqDcdjn1werY2ur_68qkjVpNhMpwqzQBGzjQ,15521
+returnn/frontend/_native/module.cpp,sha256=lS1Oypo3n6oCu6cxKAmqpNjSvQN9aMZIOeMec96FWYU,15626
 returnn/frontend/_native/module.hpp,sha256=uf4HPSTrFP2brGR_x9G5N1ZlZ-ok5GakMbNo4LbqxUg,6670
 returnn/frontend/_native/py_utils.hpp,sha256=vcxKGmOyDRuwsmmSEjoaCJyKMy1BNYoGlso2pZu7VoE,3139
-returnn/frontend/_native/tensor_ops.cpp,sha256=G_geJq_9yaTYoG7YsZydiTTkJoqdYB-D8CF8g-uaiec,70068
+returnn/frontend/_native/tensor_ops.cpp,sha256=bYtwwn_NeJfAEHWYPEJlkoLDKt9baZ3RA8av7gtz2qc,70246
 returnn/frontend/_native/tensor_ops.hpp,sha256=dDqvUejRNHjItnmOP5aHyAQbAmXmXoDVXSe3tveEU8A,3732
 returnn/frontend/audio/__init__.py,sha256=8mahwucBje8qHKw0bOvoySlvvD0rFKxviSvcAHSjiJY,67
 returnn/frontend/audio/mel.py,sha256=VZdxf2mTLzLOXsLRzCvaad712Zf0c2iwdthrzeVfgxk,7885
@@ -177,7 +177,7 @@ returnn/tf/sprint.py,sha256=Yqjh0-6sCWHpdDPQCzHKx7TwQCOjJyjfd0KHtnYdd-8,5471
 returnn/tf/updater.py,sha256=St4Z5iBjlkWaB6CiS-K1VNc_iLaan2e6-mVMTTPldzk,72034
 returnn/tf/frontend_layers/README.md,sha256=P4vVl_EK-4jT55m40mq-K4Nr9yFY0tJR5fmDzTHSDFE,1096
 returnn/tf/frontend_layers/__init__.py,sha256=MGUn7rv6fOefbtkX-5pq6fC1T6Y5h0oh1uOPSEcv1_I,506
-returnn/tf/frontend_layers/_backend.py,sha256=AvmukcCbuUqsOiZz9OmpF9oYvTY_XT-kSsY0kD-neyg,47328
+returnn/tf/frontend_layers/_backend.py,sha256=8lWE6LxxdNx8FnFvp2Pnk-UqJ8oymxXwx7s9HTEgDug,47443
 returnn/tf/frontend_layers/_utils.py,sha256=ijByaDOqPDod5mZC9EoTkt8PHBEODXHsWbkwDOF9XW4,4205
 returnn/tf/frontend_layers/cond.py,sha256=yQ2h5W0sgMZndJdrWv2EE9k9yIcspQ1U0HwBSh3hOKE,14830
 returnn/tf/frontend_layers/config_entry_points.py,sha256=t01RWOiaZohzuqPXX-MLV0P5yCOfE0dz-9dZ77_pK4c,5751
@@ -190,7 +190,7 @@ returnn/tf/frontend_layers/masked_computation.py,sha256=I_TW0Qm4Yl_wPZ6TkuK7a-wB
 returnn/tf/frontend_layers/parameter_assign.py,sha256=B_7kgobRyFtExiuSy2MsVGpAR36-jdG-xKABGc6EUGM,5103
 returnn/tf/frontend_layers/prev_tensor_ref.py,sha256=EqTAanOgYAhl8o2fMylN52mfReH9heAQFdzn9CwqAX4,2282
 returnn/tf/frontend_low_level/__init__.py,sha256=34469k3KzMUIGowxReOZnbf6WdTjxY73Gp1a4WqDN1M,62
-returnn/tf/frontend_low_level/_backend.py,sha256=C__svMpechuQq9Q2mEAEVhy8vQPGfezhNttsI1K-sVE,24541
+returnn/tf/frontend_low_level/_backend.py,sha256=JwwRRIGnElqBC4bTImdB7w3U1u_SJESeZHYLmq86wog,24479
 returnn/tf/layers/__init__.py,sha256=Ngu-X84nWFgz7ndDu88DqoZ-5lUMMTQWH4g7N8pSoCg,72
 returnn/tf/layers/base.py,sha256=KcADpZUxqLkoFpQPMe_l9thRC7rpyBJIZCHITmnOd7M,153169
 returnn/tf/layers/basic.py,sha256=IVQ_6PkM-uuBN_vVg-VeGM74bb1pc6TjJhKf92pPS1I,610870
@@ -206,7 +206,7 @@ returnn/tf/util/ken_lm.py,sha256=R60UAoywriuDIeQ2Hk3Vm_waf2Hxxc88ofzEw6X6Sd4,173
 returnn/tf/util/open_fst.py,sha256=sZRDw4TbxvhGqpGdUJWy1ebvlZm4_RPhygpRw9uLAOQ,11265
 returnn/torch/README.md,sha256=jzJ2FpOHW02vxN69yKaV97C9LI-hmvjBglKfdZXIDdc,85
 returnn/torch/__init__.py,sha256=MHEUyNHB20Vy89uKAqZoj6FxJKF1Gq3HW-i6ra1pNcI,24
-returnn/torch/distributed.py,sha256=i13cUVjI7GxpO0TAresrNyCM0ZBAaf-cXNr09Fmg_2k,6266
+returnn/torch/distributed.py,sha256=skFyutdVztxgTEk3HHJ8S83qRWbNpkNT8Tj16Ic0_hE,6981
 returnn/torch/engine.py,sha256=8BIpdcrpbJL9HrvCX-hISh-14zW9aSrHGvRWT9s0zOk,77103
 returnn/torch/updater.py,sha256=GqtBvZpElPVMm0lq84JPl4NVLFFETZAzAbR0rTomSao,28249
 returnn/torch/data/__init__.py,sha256=6cLNEi8KoGI12PF6akN7mI_mtjlx-0hcQAfMYoExwik,132
@@ -216,7 +216,7 @@ returnn/torch/data/queued_data_iter.py,sha256=PoOsGHdHVZjTmcyfq_ZOw--P6hyfTdmAWI
 returnn/torch/data/returnn_dataset_wrapper.py,sha256=1Bw82-Ge_8m_DSDXZNqQ3zGDic2HQlp6jysELL0NVK0,7369
 returnn/torch/data/tensor_utils.py,sha256=-Teqi--LLbt6q_5mDRdoHZHmPgSdC83W706ukif_YiU,1284
 returnn/torch/frontend/__init__.py,sha256=AA48HZnC17ASuKA0EWy8loZ-Bib_yUtqF4T1wYvjst4,62
-returnn/torch/frontend/_backend.py,sha256=sdC4kj7_2wyEVxX536TunbY4VbFaScSnMIBi1F27bsE,100128
+returnn/torch/frontend/_backend.py,sha256=mjR6Ilt2zlnIO4_CpVPCLQ0XVJa_QmW3HsZtR2KT8yk,101110
 returnn/torch/frontend/_rand.py,sha256=1JgIkV2XmpgJD86zXZ-NCAe-QuoP2swr6NaS1oz3Qa8,1830
 returnn/torch/frontend/bridge.py,sha256=Z2_UW8AagezC7zsXDc5PKcd8G9WwisV7j9SWGHU0m4U,7840
 returnn/torch/frontend/raw_ops.py,sha256=lF0h-KtYYsdaaqQADylVZp9qzPskOOXA4MfmYDyx5IU,296
@@ -233,7 +233,7 @@ returnn/torch/util/gradient_checkpoint.py,sha256=iLy-FB65DC8O6LxzmMvFjnSdpIVpko8
 returnn/torch/util/module.py,sha256=MXHIrF9Isu575DDJIa81212ULKwdqu1oOLxDVZecVSk,1693
 returnn/torch/util/scaled_gradient.py,sha256=3585VuNypBty-pW6r3BKK047H3MqZQSdMjXeYAb4cmU,3192
 returnn/util/__init__.py,sha256=UIG1qw4idqhW71BV60ha7h9PktxvEVcBIu0lYRossK8,336
-returnn/util/basic.py,sha256=nhCfxWwGL7FchgFW5x9V2OgXD0HtpN885NASdwfeKYg,142339
+returnn/util/basic.py,sha256=__rtDp8crZfm0mEeAKsRxNCdWuBHh9OeOm8UO-X4CJU,142380
 returnn/util/better_exchook.py,sha256=MVMnuu6KoyqgvlMeQLQNTfdspcPR9MwigCXOpeTVqCI,62956
 returnn/util/bpe.py,sha256=LWFhICZsEOnMwNws0lybPNzKRX6rSr8yKCvP65vjl9Y,19656
 returnn/util/debug.py,sha256=wuRzdg9zB84WWCGyTjmRR_zYypu8gXxlc0nZ6si9OC8,28224
@@ -253,8 +253,8 @@ returnn/util/sig_proc.py,sha256=Tjz0VOAVyqu2qDCF5HZ1JjALjcFsHcNkcd96WgZeKfE,7265
 returnn/util/task_system.py,sha256=y4sMVXQ25Qd2z0rx03uOlXlkE-jbCYC1Sjfn-XlraVU,26003
 returnn/util/train_proc_manager.py,sha256=Pjht28k6uz6BNQ47uW6Gf880iyq5q4wx7P_K2tmoAM8,3266
 returnn/util/watch_memory.py,sha256=BR5P2kvBN6UI81cE0_1WAA6Hd1SByLbBaiDxvLhPOew,4213
-returnn-1.20250220.174943.dist-info/LICENSE,sha256=ywBD_U2aD4vpuoIgNAsjIGBYydl0tVKll3De0Z8s77c,11041
-returnn-1.20250220.174943.dist-info/METADATA,sha256=SqIvwRtUoU4A4DTo3MA0lH5LDudxkcTNX0sV2dn-cNM,5215
-returnn-1.20250220.174943.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-returnn-1.20250220.174943.dist-info/top_level.txt,sha256=Lsn4WZc5Pbfk0-xDQOgnFCxOoqxL4CyeM3N1TFbJncw,8
-returnn-1.20250220.174943.dist-info/RECORD,,
+returnn-1.20250221.114352.dist-info/LICENSE,sha256=ywBD_U2aD4vpuoIgNAsjIGBYydl0tVKll3De0Z8s77c,11041
+returnn-1.20250221.114352.dist-info/METADATA,sha256=VL6JCmrRQYGi4OdodEygiq3hJyM1C-73zvMTSusYpOs,5215
+returnn-1.20250221.114352.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+returnn-1.20250221.114352.dist-info/top_level.txt,sha256=Lsn4WZc5Pbfk0-xDQOgnFCxOoqxL4CyeM3N1TFbJncw,8
+returnn-1.20250221.114352.dist-info/RECORD,,