returnn 1.20251027.224345__py3-none-any.whl → 1.20260109.93428__py3-none-any.whl

returnn/torch/frontend/_backend.py

@@ -1166,20 +1166,29 @@ class TorchBackend(Backend[torch.Tensor]):
         if start is None:
             start = 0
         if isinstance(size, Dim):
+            assert end is None
             size = size.get_dim_value()
         elif isinstance(size, Tensor):
+            assert end is None
             assert size.dims == ()  # scalar
             size = size.raw_tensor
-        if size is not None:
-            assert end is None
-            out.raw_tensor = torch.narrow(source.raw_tensor, dim=axis_int, start=start, length=size)
-        else:
+        elif isinstance(size, int):
+            pass
+        elif size is None:
             if isinstance(end, Tensor):
                 assert end.dims == ()
                 end = end.raw_tensor
-            if end is None:
+            elif isinstance(end, int):
+                if end < 0:
+                    end += axis.get_dim_value()
+            elif end is None:
                 end = axis.get_dim_value()
-            out.raw_tensor = torch.narrow(source.raw_tensor, dim=axis_int, start=start, length=end - start)
+            else:
+                raise TypeError(f"slice: unsupported type for end: {type(end)}")
+            size = end - start
+        else:
+            raise TypeError(f"slice: unsupported type for size: {type(size)}")
+        out.raw_tensor = torch.narrow(source.raw_tensor, dim=axis_int, start=start, length=size)
         return out
 
     @staticmethod
@@ -1352,12 +1361,24 @@ class TorchBackend(Backend[torch.Tensor]):
         a_dims = a.dims
         b_dims = b.dims
 
-        assert all(dim in a_dims for dim in reduce), (
-            f"'a' does not have the specified reduce dim(s) {reduce} (a dims: {a_dims})"
-        )
-        assert all(dim in b_dims for dim in reduce), (
-            f"'b' does not have the specified reduce dim(s) {reduce} (b dims: {b_dims})"
-        )
+        if not all(dim in a_dims for dim in reduce) or not all(dim in b_dims for dim in reduce):
+            # revert to the generic einsum implementation
+            assert all(dim in a_dims + b_dims for dim in reduce), "Some reduce Dims not in a or b."
+            result_dims = [dim for dim in a_dims if dim not in reduce] + [
+                dim for dim in b_dims if dim not in reduce and dim not in a_dims
+            ]
+            map_to_letter = {}
+            for dim in a_dims + b_dims:
+                if dim not in map_to_letter:
+                    map_to_letter[dim] = chr(97 + len(map_to_letter))  # 'a', 'b', 'c', ...
+            a_subscript = "".join(map_to_letter[dim] for dim in a_dims)
+            b_subscript = "".join(map_to_letter[dim] for dim in b_dims)
+            out_subscript = "".join(map_to_letter[dim] for dim in result_dims)
+            raw_result = torch.einsum(f"{a_subscript},{b_subscript}->{out_subscript}", a.raw_tensor, b.raw_tensor)
+            result_tensor = Tensor(
+                "einsum", dims=result_dims, raw_tensor=raw_result, dtype=TorchBackend.get_dtype_name_raw(raw_result)
+            )
+            return result_tensor
 
         if len(reduce) > 1:
             reduce = list(reduce)
@@ -1767,6 +1788,9 @@ class TorchBackend(Backend[torch.Tensor]):
         remaining_dims = [d for d in tensor.dims if d not in mask.dims]
         tensor_templ_dims = tuple(dims) + tuple(remaining_dims)
         in_raw = tensor.copy_compatible_to_dims_raw(tensor_templ_dims)
+        if any(in_raw.shape[i] == 1 < d.get_dim_value() for i, d in enumerate(dims)):
+            # unbroadcast
+            in_raw = in_raw.expand([d.get_dim_value() for d in tensor_templ_dims])
         if mask.raw_tensor.device.type == "meta":
             # This is not supported, but also, we would anyway not know the out shape.
             # However, instead of erroring, just assume some dummy mask.
@@ -1920,7 +1944,7 @@ class TorchBackend(Backend[torch.Tensor]):
         if not out_spatial_dims:
             out_spatial_dims = rf.make_conv_out_spatial_dims(
                 in_spatial_dims=in_spatial_dims,
-                filter_size=[d.dimension for d in filter_size],
+                filter_size=filter_size,
                 strides=strides or 1,
                 dilation_rate=dilation_rate or 1,
                 padding=padding,
@@ -2033,6 +2057,104 @@ class TorchBackend(Backend[torch.Tensor]):
         out.feature_dim = out_dim
         return out, out_spatial_dims
 
+    # noinspection PyShadowingBuiltins
+    @staticmethod
+    def transposed_conv(
+        source: Tensor,
+        *,
+        in_dim: Dim,
+        out_dim: Dim,
+        in_spatial_dims: Sequence[Dim],
+        out_spatial_dims: Optional[Sequence[Dim]] = None,
+        filter: Tensor,
+        filter_size: Sequence[Dim],
+        padding: str,
+        remove_padding: Union[Sequence[int], int] = 0,
+        output_padding: Optional[Union[Sequence[Optional[int]], int]] = None,
+        strides: Optional[Sequence[int]] = None,
+        bias: Optional[Tensor] = None,
+    ) -> Tuple[Tensor, Sequence[Dim]]:
+        """transposed convolution"""
+        if not out_spatial_dims:
+            out_spatial_dims = rf.make_transposed_conv_out_spatial_dims(
+                in_spatial_dims=in_spatial_dims,
+                filter_size=filter_size,
+                strides=strides,
+                padding=padding,
+                output_padding=output_padding,
+            )
+            assert remove_padding == 0  # not implemented yet otherwise...
+        if strides is None:
+            strides = [fs.dimension for fs in filter_size]
+        filter_dims = (in_dim, out_dim) + tuple(filter_size)
+        filter = filter.copy_transpose(filter_dims)
+        batch_dims = [d for d in source.dims if d not in (in_dim,) + tuple(in_spatial_dims)]
+        # Torch conv expects (N,C,<spatial dims>) as shape.
+        source = source.copy_transpose(batch_dims + [in_dim] + list(in_spatial_dims))
+        if len(batch_dims) == 1:
+            src_raw = source.raw_tensor
+        else:
+            src_raw = torch.reshape(
+                source.raw_tensor,
+                # potentially merge batch dims all together
+                [-1, in_dim.get_dim_value()] + [d.get_dim_value() for d in in_spatial_dims],
+            )
+        if padding == "same":
+            raise NotImplementedError("transposed_conv with padding='same' not implemented")
+        if padding == "valid":
+            padding_val = 0
+        else:
+            raise ValueError(f"invalid padding {padding!r}, expected 'same' or 'valid'")
+        if len(filter_size) == 1:
+            out_raw = torch.nn.functional.conv_transpose1d(
+                src_raw,
+                weight=filter.raw_tensor,
+                bias=bias.raw_tensor if bias is not None else None,
+                stride=strides,
+                padding=padding_val,
+                output_padding=output_padding or 0,
+            )
+        elif len(filter_size) == 2:
+            out_raw = torch.nn.functional.conv_transpose2d(
+                src_raw,
+                weight=filter.raw_tensor,
+                bias=bias.raw_tensor if bias is not None else None,
+                stride=strides,
+                padding=padding_val,
+                output_padding=output_padding or 0,
+            )
+        elif len(filter_size) == 3:
+            out_raw = torch.nn.functional.conv_transpose3d(
+                src_raw,
+                weight=filter.raw_tensor,
+                bias=bias.raw_tensor if bias is not None else None,
+                stride=strides,
+                padding=padding_val,
+                output_padding=output_padding or 0,
+            )
+        else:
+            raise ValueError(f"invalid number of filter dims {filter_size}, expected 1, 2, or 3")
+        if remove_padding:
+            if isinstance(remove_padding, int):
+                remove_padding = [remove_padding] * len(out_spatial_dims)
+            assert len(remove_padding) == len(out_spatial_dims)
+            slices = [slice(None)] * out_raw.ndim
+            for i, pad in enumerate(remove_padding):
+                if pad > 0:
+                    slices[2 + i] = slice(0, -pad)
+            out_raw = out_raw[tuple(slices)]
+        out = Tensor(
+            "transposed_conv",
+            dims=batch_dims + [out_dim] + list(out_spatial_dims),
+            dtype=TorchBackend.get_dtype_name_raw(out_raw),
+        )
+        if len(batch_dims) == 1:
+            out.raw_tensor = out_raw
+        else:
+            out.raw_tensor = torch.reshape(out_raw, [d.get_dim_value() for d in out.dims])
+        out.feature_dim = out_dim
+        return out, out_spatial_dims
+
     @staticmethod
     def pool(
         source: Tensor,

returnn/torch/frontend/bridge.py

@@ -136,6 +136,15 @@ class RFModuleAsPTModule(torch.nn.Module):
     def _get_name(self):
         return self._rf_module.__class__.__name__ + "[RF→PT]"
 
+    def __repr__(self) -> str:
+        """
+        Return a custom repr for Sequential/ModuleList that compresses repeated module representations if possible,
+        otherwise fallback to default behavior.
+        """
+        if _can_use_compact_repr(self):
+            return _repr_compact(self)
+        return super().__repr__()
+
     @property
     def rf_module(self) -> rf.Module:
         """RF module"""
@@ -193,3 +202,55 @@ class RFModuleAsPTModule(torch.nn.Module):
             # See similar logic in torch.nn.Module._apply.
             pt_param = torch.nn.Parameter(tensor, tensor.requires_grad)
             rf_param.raw_tensor = pt_param
+
+
+def _can_use_compact_repr(self: RFModuleAsPTModule) -> bool:
+    return list(self._modules.keys()) == [str(i) for i in range(len(self._modules))]
+
+
+def _repr_compact(self: RFModuleAsPTModule) -> str:
+    """
+    Return a custom repr for Sequential/ModuleList that compresses repeated module representations.
+    Code copied and adapted from torch.nn.ModuleList.__repr__.
+    """
+    list_of_reprs = [repr(item) for item in self._modules.values()]
+    if len(list_of_reprs) == 0:
+        return self._get_name() + "()"
+
+    start_end_indices = [[0, 0]]
+    repeated_blocks = [list_of_reprs[0]]
+    for i, r in enumerate(list_of_reprs[1:], 1):
+        if r == repeated_blocks[-1]:
+            start_end_indices[-1][1] += 1
+            continue
+
+        start_end_indices.append([i, i])
+        repeated_blocks.append(r)
+
+    lines = []
+    main_str = self._get_name() + "("
+    for (start_id, end_id), b in zip(start_end_indices, repeated_blocks):
+        local_repr = f"({start_id}): {b}"  # default repr
+
+        if start_id != end_id:
+            n = end_id - start_id + 1
+            local_repr = f"({start_id}-{end_id}): {n} x {b}"
+
+        local_repr = _add_indent(local_repr, 2)
+        lines.append(local_repr)
+
+    main_str += "\n  " + "\n  ".join(lines) + "\n"
+    main_str += ")"
+    return main_str
+
+
+def _add_indent(s_: str, num_spaces: int) -> str:
+    s = s_.split("\n")
+    # don't do anything for single-line stuff
+    if len(s) == 1:
+        return s_
+    first = s.pop(0)
+    s = [(num_spaces * " ") + line for line in s]
+    s = "\n".join(s)
+    s = first + "\n" + s
+    return s

returnn/torch/util/exception_helper.py

@@ -71,7 +71,13 @@ def help_on_torch_exception(
     if not count_frames:
         exc_ext.append("(No module call frames.)")
 
-    if len(exc.args) == 1 and isinstance(exc.args[0], str) and not always_direct_print:
+    if (
+        # KeyError formatting would be wrong, showing `KeyError: "enc_spatial_dim\n\nStep idx: 0\..."`
+        not isinstance(exc, KeyError)
+        and len(exc.args) == 1
+        and isinstance(exc.args[0], str)
+        and not always_direct_print
+    ):
         exc.args = ("\n".join([exc.args[0], ""] + exc_ext),)
     else:
         for msg in exc_ext:

returnn/util/basic.py

@@ -365,12 +365,9 @@ def get_checkpoint_filepattern(filepath):
     :return: CheckpointLoader compatible filepattern
     :rtype: str
     """
-    if filepath.endswith(".meta"):
-        return filepath[: -len(".meta")]
-    elif filepath.endswith(".index"):
-        return filepath[: -len(".index")]
-    elif filepath.endswith(".pt"):
-        return filepath[: -len(".pt")]
+    for ext in [".meta", ".index", ".pt"]:
+        if filepath.endswith(ext):
+            return filepath[: -len(ext)]
     return filepath
 
 

returnn/util/better_exchook.py

@@ -1093,6 +1093,7 @@ def format_tb(
     with_color=None,
     with_vars=None,
     clear_frames=True,
+    colorize=None,
 ):
     """
     Formats a traceback into a list of strings, each corresponding to one frame.
@@ -1110,11 +1111,14 @@ def format_tb(
         That will potentially fix some mem leaks regarding locals, so it can be important.
         Also see https://github.com/python/cpython/issues/113939.
         However, any further access to frame locals will not work (e.g., if you want to use a debugger afterward).
+    :param colorize: for compat with Python >=3.13, currently ignored
     :return: list of strings, each corresponding to one frame in the traceback.
         Each string contains the file name, line number, function name, source code line, maybe relevant variables,
         etc., and a final newline.
     :rtype: list[str]
     """
+    if colorize is not None and with_color is None:
+        with_color = colorize
     color = Color(enable=with_color)
     output = _OutputLinesCollector(color=color)
 

returnn/util/debug.py

@@ -704,7 +704,7 @@ def check_py_traces_rf_to_pt_equal(
     """
     import random
     import torch
-    from returnn.tensor import Tensor, Dim
+    from returnn.tensor import Dim
     import returnn.frontend as rf
 
     # noinspection PyProtectedMember
@@ -715,9 +715,18 @@ def check_py_traces_rf_to_pt_equal(
     def _get_entry(trace, func, i, name, j):
         return trace[func][i][name][j]
 
+    def _get_entry_attr(trace, func, i, name, j):
+        name, attr = name.split(".", 1)
+        obj = trace[func][i][name][j]
+        return eval(f"{name}.{attr}", {name: obj})
+
     def _resolve_dim(dim: Union[Dim, str]) -> Dim:
         if isinstance(dim, Dim):
             return dim
+        elif isinstance(dim, str) and "." in dim:
+            dim = _get_entry_attr(trace_rf, *check_rf[:2], dim, -1)
+            assert isinstance(dim, Dim)
+            return dim
         elif isinstance(dim, str):
             dim = _get_entry(trace_rf, *check_rf[:2], dim, -1)
             assert isinstance(dim, Dim)
@@ -763,7 +772,7 @@ def check_py_traces_rf_to_pt_equal(
             if len(indices) > 5:
                 msgs.append("  non-matching ...")
             non_matching.append("\n".join(msgs_prefix + msgs))
-            print(f"  mismatch!")
+            print("  mismatch!")
             for msg in msgs:
                 print(msg)
 

returnn/util/file_cache.py

@@ -426,7 +426,21 @@ class FileCache:
                 orig_mtime_ns = os.stat(src_filename).st_mtime_ns
                 FileInfo(mtime_ns=orig_mtime_ns).save(info_file_name)
 
-                _copy_with_prealloc(src_filename, dst_tmp_filename)
+                try:
+                    _copy_with_prealloc(src_filename, dst_tmp_filename)
+                except Exception:
+                    # Cleanup if it was created already.
+                    # That avoids some of the ambiguity of the existence of the .copy file.
+                    # https://github.com/rwth-i6/returnn/issues/1785
+                    try:
+                        os.remove(dst_tmp_filename)
+                    except FileNotFoundError:
+                        pass
+                    try:
+                        os.remove(info_file_name)
+                    except FileNotFoundError:  # not really expected here, but safe to ignore
+                        pass
+                    raise
                 os.rename(dst_tmp_filename, dst_filename)
 
     @staticmethod

returnn/util/task_system.py

@@ -671,7 +671,7 @@ class Pickler(_BasePickler):
                 return
         # For some reason, Numpy fromstring/tostring is faster than Numpy loads/dumps.
         self.save(make_numpy_ndarray_fromstring)
-        self.save((obj.tostring(), str(obj.dtype), obj.shape))
+        self.save((obj.tobytes(), str(obj.dtype), obj.shape))
         self.write(pickle.REDUCE)
 
     dispatch[numpy.ndarray] = save_ndarray

{returnn-1.20251027.224345.dist-info → returnn-1.20260109.93428.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20251027.224345
+Version: 1.20260109.93428
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer
@@ -36,7 +36,7 @@ Welcome to RETURNN
 `RETURNN paper 2018 <https://arxiv.org/abs/1805.05225>`_.
 
 RETURNN - RWTH extensible training framework for universal recurrent neural networks,
-is a Theano/TensorFlow-based implementation of modern recurrent neural network architectures.
+is a PyTorch/TensorFlow-based implementation of modern recurrent neural network architectures.
 It is optimized for fast and reliable training of recurrent neural networks in a multi-GPU environment.
 
 The high-level features and goals of RETURNN are:

{returnn-1.20251027.224345.dist-info → returnn-1.20260109.93428.dist-info}/RECORD

@@ -1,14 +1,14 @@
-returnn/PKG-INFO,sha256=GJaNdcwgbvQQFiUHTgzu_yvcMXGo3HTUjilKYreP16c,5215
+returnn/PKG-INFO,sha256=8G2OFR-V5IlE98f0vmLneA27jg9-B7eN973G7vJpj0I,5215
 returnn/__init__.py,sha256=biBtRsM0WZ406vShaeH-9WFoqJ8XwTbn6g0EeFJ7l8E,1012
 returnn/__main__.py,sha256=lHyZcu_0yc9f7Vf_Kfdy9PmeU0T76XVXnpalHi5WKro,31740
 returnn/__old_mod_loader__.py,sha256=nvsNY-xELdS_IPNkv66Q9Rmvg4dbGW0-EBRDcCmctos,7654
 returnn/__setup__.py,sha256=22kQn2fh11iPM0hLb2Fy5sLmoU1JGvmDxXRYuRgQkwU,4659
-returnn/_setup_info_generated.py,sha256=K1u1U6k-3ADEGhb49QEaG0HItc4bK75ha5sSatjbAB4,77
-returnn/config.py,sha256=3tmKhB6FnQZaNdtcYsiB61JnEY--iZ2qmJ4yq0b6tE0,29140
+returnn/_setup_info_generated.py,sha256=RHjC4xFQRTza5prANYrhwttWMqAEZoLiKwzMpCmll80,77
+returnn/config.py,sha256=JK8EjDsUdyY2c90s0KY1rLD1kesVfz6vRT0gxy_AQ5I,29142
 returnn/forward_iface.py,sha256=A_OJiaXsX4MlXQRzST86ylyxSUZbC402PQL1REcqHjM,911
 returnn/learning_rate_control.py,sha256=ZvWryAn_tv9DhV8sh1LV3eE34Yltl3On3mYZAG4hR9s,34684
 returnn/log.py,sha256=WoTDv4XDovgvgXa7iiav-nA8pb25lOEzndbnVrDLfUo,12319
-returnn/native_op.cpp,sha256=KcBoq875Jc7P9o77P_ZY0CLg_gCnbHV4KLksI_C5Yas,35760
+returnn/native_op.cpp,sha256=itGDV05Nqg59g37qXAt-Z4c-rX7eDRQFY9efDPjaVlg,38613
 returnn/native_op.py,sha256=4_NnvfNxsM8GE_FsD6yOg6PZegqIdtJ3Sl1GdBWmFvg,244424
 returnn/pretrain.py,sha256=MHiXJZqkQFmDVyaYsGpd_Acv20wxl7Pr6s6qJzAT2FI,22648
 returnn/datasets/__init__.py,sha256=PvDlfDOaaopIeUIt0OSvHD2eHZkdkyE-sjMXf35EH5U,390
@@ -21,13 +21,13 @@ returnn/datasets/distrib_files.py,sha256=48edqdf7YpnPJ-TOis3Mz5U9A2DSxfiYT1HCMSt
 returnn/datasets/generating.py,sha256=o9-JZ2s5QKssux6GcSaM3oivf_PE6nhSOeytRyGB7pQ,99574
 returnn/datasets/hdf.py,sha256=v5sjBenURR9Z-g7AQ9tsL84yDSye5RtbLpym3M6HSDE,67833
 returnn/datasets/huggingface.py,sha256=ls9WMR6gUcMgGksl80g0An1az5Xjya_V3ojbbbsZqrU,20047
-returnn/datasets/lm.py,sha256=rQ3jV43lSnlGkKu7m5jTTH7aK0BOMXQocsHfJ8OGec8,99950
+returnn/datasets/lm.py,sha256=CXl_g-Z28RWlBTzx35uC4r_GCwOP05LIsUp0iSi6JG4,100652
 returnn/datasets/map.py,sha256=kOBJVZmwDhLsOplzDNByIfa0NRSUaMo2Lsy36lBvxrM,10907
-returnn/datasets/meta.py,sha256=VJ5bk8esq2-b9likNSrCsHQKiLC3Vvti5oBAxg-AsIk,99422
+returnn/datasets/meta.py,sha256=hTtfwINIxP2S4JQ5IQXzvTh2MixwxzeF06pPTW36yl0,101456
 returnn/datasets/multi_proc.py,sha256=BClXq0fActi1XQa4vcMhHmhYF0Q-fnnDzlIlbBM6_DM,22614
 returnn/datasets/normalization_data.py,sha256=J3njQCMvWAbIAVPepO2L_Xdau9eWYB7Zyd6STeGzTbc,14615
 returnn/datasets/numpy_dump.py,sha256=wl8bKIKAlff2HPJPtuu5wBg3TLOf16d2wLVB4lLAwTM,5158
-returnn/datasets/postprocessing.py,sha256=6SfT58BxbHYO2QlGzOgIV04Zqkp-kl0B85168DQaB9A,24060
+returnn/datasets/postprocessing.py,sha256=Yy7rSzpB8z6PkUTZsPE_AN9Di8FPdNn617JTCV7L-VI,42453
 returnn/datasets/raw_wav.py,sha256=M7eTHp4CTtLQf3yPTiJY-mSJYgZNxkGV9IFN9J1dq_4,9144
 returnn/datasets/sprint.py,sha256=JAs5dOmdteSOwA7YQcTF9KaTCtGfRjiyJUZClSr85pY,55502
 returnn/datasets/stereo.py,sha256=PkowC91bZWihIYuIZgyGgPcNwgq5jBvyxxu1nER-VhM,17633
@@ -35,7 +35,7 @@ returnn/datasets/text_dict.py,sha256=xOWwZc5xGVg2Ic1Ezbm-uEcj_i7ajxqRfPn_TihQbzc
 returnn/datasets/util/__init__.py,sha256=rEKhSD6fyhDiQF-x7dUQMwa29JZu72SDm7mYcCcLghY,52
 returnn/datasets/util/feature_extraction.py,sha256=axtXDb9wcNpOmyhmW3WJUj5xda29TKkKvOcGGvq7ExA,23923
 returnn/datasets/util/strings.py,sha256=pP8pmXhArkssYqmPOLuxEG9gsko891ZxrWiai86qbLE,412
-returnn/datasets/util/vocabulary.py,sha256=1W13FgxPVP9XSIyhkt4I7CXK5lj99zT97R-gLa_xnAU,27964
+returnn/datasets/util/vocabulary.py,sha256=994cHmRI3Yy8mHL79oCMrJITRIS9su11V2zizF5__Fo,31389
 returnn/engine/__init__.py,sha256=br7hpn8i_hIBi2uTQfnN3BF9g5DREYa_mQi0_Nvlu6o,228
 returnn/engine/base.py,sha256=0n4FtB_B2H3W_9KdoLr0P7YPER-hVkbk69pwFqsqmqw,18467
 returnn/engine/batch.py,sha256=amXW8mGspuSQjo00JdisE2eOLy5Ij1weWWzkE-lXSJM,9912
@@ -81,16 +81,16 @@ returnn/frontend/_cache.py,sha256=Uao2xzfvVaKABk1fkxcpXzxKIGJaI9FwwlTvvoNUstk,85
 returnn/frontend/_numpy_backend.py,sha256=fZjks7p3dgxVZ6tSDazTTgBxNjJqXjfqgw_7mA7rDEE,9066
 returnn/frontend/_random_journal.py,sha256=_ktP_mjgx8vtQQGX_DofdhewJj0aPiczefTWeemPkmo,5457
 returnn/frontend/_utils.py,sha256=uVQldGHyYKIyhSEmumJ04ix5eP5tjZw4CEC0w6-zhyQ,12074
-returnn/frontend/array_.py,sha256=ci9NnYqwDxryOoiHCNg8DbOb9yWJScVSm7FnCqgywPY,54257
-returnn/frontend/attention.py,sha256=GKt-Xqnz8sIyXVrE0i4VCS7J2Wu7dmoH_BA0Cu8CrXQ,45769
+returnn/frontend/array_.py,sha256=bZwTgNkMsGiSP6TVgI7bxY6zZMjcs9TVsHlajYrHUoA,56791
+returnn/frontend/attention.py,sha256=bFD9Ei6GxSi-BC1OfueDyTIE-51a3dKKZOWdSIbz7l8,46633
 returnn/frontend/backend.py,sha256=iQ9w4xl8Ea7bgpb0VUaCKq50rV5Bl2E5J8Rhd-oqD_c,883
 returnn/frontend/build_from_dict.py,sha256=rfWa2rjjhIR_kIQED_nMrygrQBunS6unegzWTLVbC98,3017
 returnn/frontend/cond.py,sha256=gh6wg0aSbAJQfKRv4BQAu-EfPWtWPLFjgc8IaPPFmwg,1023
 returnn/frontend/const.py,sha256=A5fP9w6Akv56d89pPvdoZaXvC9ZTYcexepnS9O2clOc,3945
 returnn/frontend/container.py,sha256=wF3OlQN7WlOVmmdapUth_Unha3DVf6h1B7okBJAuJDA,8011
 returnn/frontend/control_flow_ctx.py,sha256=v17CsNwRnZYe8GdMtGJt2ftibfxMCGK1i0l-GX5ILu0,699
-returnn/frontend/conv.py,sha256=4Mrq7MFc0f7SJ8g-wJEv4Lg3Stmju-fMwD09qKv6CuQ,32174
-returnn/frontend/device.py,sha256=Sjara0EmFLhu9O55cN_p6OwU0NgdNCCQjyAuQhiWpGw,1437
+returnn/frontend/conv.py,sha256=RbVyFGspn40VNT1B-KWWaDBBUhd7VFhKTN-V_SrwPlU,39514
+returnn/frontend/device.py,sha256=gX7zPRZrnhjMgpgg6aACtE7Lg6qYlzerYbeCTiRhxhw,1665
 returnn/frontend/dims.py,sha256=_HDU-Kxn3pApicFkm0F4Fs-ZAuF1gKXG8rroQHCFQQI,13073
 returnn/frontend/dropout.py,sha256=TjqZCKDIOBeHr14-NCemOm9m3p84LxQuPH1DvRAYg88,5028
 returnn/frontend/dtype.py,sha256=Ooc5BrcNrTp6XShuFEV9g5V6-niuy4ImP_Lt_Qgq3jE,1886
@@ -101,8 +101,8 @@ returnn/frontend/init.py,sha256=bVB7bpghaY8DI_HL0mkB_9z95onWnIX2zlW4hlMYnRw,7494
 returnn/frontend/label_smoothing.py,sha256=lxmaowNr61sCMzMewqHhu1r0CcklYfhLXlFnBu8DeAU,5676
 returnn/frontend/linear.py,sha256=xRUjnkD3MTWDezSaYATBYJQ2fa1RhKMNrTuhC54hhVs,2252
 returnn/frontend/loop.py,sha256=t-z6ke1X03I2aPUEqLYmVZWyMzfW3IedFvKUGc-TCX8,16160
-returnn/frontend/loss.py,sha256=uSvou2MPd13JiLAg_OIQ3AyyLvD3RHjMEVgFEN0gKqU,7440
-returnn/frontend/math_.py,sha256=A_RkZ5lH2uXMchfPIH3itraWtMNNCVckQHHpf7aIIZQ,17295
+returnn/frontend/loss.py,sha256=aSKzjhjIikeNJqzcUBBlaBTXILuAoW6wrmsExGtJJBY,8572
+returnn/frontend/math_.py,sha256=wIWYtjcIEV_QXNJiNT1lYsVRQdNLxtBbTYiAPr_OR3Y,18442
 returnn/frontend/matmul.py,sha256=xkueyxzSDz8MsYaWxPSjmV2Yy-tcaiOQDXbFt1IQM2A,1944
 returnn/frontend/module.py,sha256=nt35I9xyHuH42qobLHGUFoNI5-mVieAtA36SqK6NhpY,11065
 returnn/frontend/nested.py,sha256=PKsKWHwE2SI19DjZ9vRI8q4-ywIGMK3-TTUuqdXrVlM,15592
@@ -139,16 +139,16 @@ returnn/frontend/decoder/__init__.py,sha256=A-koKyPVlXp_V_2bk6GKZ1Xfv4rYIcfxGMXQ
 returnn/frontend/decoder/transformer.py,sha256=64Z1IY_WcDuj8Ti73BGwbT_grrEpxBl5mIsBZkqJzHQ,24650
 returnn/frontend/encoder/__init__.py,sha256=0QGLlujRIKx3zBREeShza_-xhGIxj73zbd7t-g1m-ho,17
 returnn/frontend/encoder/base.py,sha256=A759EwCYAmSi-kzXz1vaTjR2l59TvNGQlzaNdp3UOKs,2109
-returnn/frontend/encoder/conformer.py,sha256=rWulygolesbYkLw9naSxwygaZhWqKpHKEVj-1AQbel0,21351
+returnn/frontend/encoder/conformer.py,sha256=I1OeaU2P7lm-N_ODS_P4BVQaplJR4Ies1Yd7Lr9mdFw,22225
 returnn/frontend/encoder/conformer_v2.py,sha256=vAYdT8m2Zzg3IIZZafeccClFHU1_c9T-EgBOsHadQPA,7701
 returnn/frontend/encoder/e_branchformer.py,sha256=SZdhpb90FaQdpzgvSOtFPLbLCa0NdycbB5Z4vMoY4TM,12279
-returnn/frontend/encoder/transformer.py,sha256=Jj0mF1D2MohOk-9sGYdsLtVW_86fwoq4pKWCdPMvPR8,11580
+returnn/frontend/encoder/transformer.py,sha256=0-ku9A8r_w3USQd0aAQ0fdPvFILNWcGaGZ7g3SE-Xjo,11656
 returnn/import_/__init__.py,sha256=L2dKxWCcn0fz_7H7OS-zw5i5Yrljjjh_d61dEcFP_JY,243
 returnn/import_/common.py,sha256=0cmvyd7NtMLH55IskEoSDtkcMwChxLhauV2UZ4mK68I,8148
 returnn/import_/git.py,sha256=IXBVOybQAHf5OlMfVY6oZ-7eiDYPG0OR7MyDJKcVHSM,13961
 returnn/import_/import_.py,sha256=q_NQRbfK5TsALakUxixE0SCqDccfGh6wkquCmJ-3s6w,798
 returnn/sprint/__init__.py,sha256=bKRS04_tJdp-z6Rmv4Fm3hGD9M3UVRcXtvcS8VG12KA,64
-returnn/sprint/cache.py,sha256=sYWQJHHvdcXDXl_SNraEsl3X8yS0F4eZjFHiuDhHkjQ,35069
+returnn/sprint/cache.py,sha256=9QxVIUoDXo2WGkenof7vuh8vqBXShf0Dr4twiLsygaQ,34976
 returnn/sprint/control.py,sha256=FFjpvoktLgp4ETxaMVKrwiegnAQLW5lD_VF4OG3ttUw,31133
 returnn/sprint/error_signals.py,sha256=FvXscTpbLWKRKsK3slPT6QfU5FytLUNblB7Ogup_L7k,23348
 returnn/sprint/extern_interface.py,sha256=l-v1X-Yg0UpTFe7Y3c4FwWOqpSNuv9Oy5EzqlKWUMlE,12055
@@ -164,7 +164,7 @@ returnn/tensor/dim.py,sha256=652DlcSe6o6l5OyY5xt9Yigij_Xry-ToG9AemMX3roY,4208
 returnn/tensor/marked_dim.py,sha256=Ae2hQIb5QixRU2gDhQEm0tmYt8TmomWoGERB414jR8o,1884
 returnn/tensor/tensor.py,sha256=IIHbDu0D_aX8U4LKTm5ThD_fuoGhn98B9EyvVBsPJ3E,9083
 returnn/tensor/tensor_dict.py,sha256=-20YPbXfRDE9WurkfQM-Mw6H8ouaBGL_90SDmK0b4cw,7534
-returnn/tensor/utils.py,sha256=B6_XyNTXPIyLxWk061Qo-Md8_DnINGdVwpXJF6pahBk,9772
+returnn/tensor/utils.py,sha256=GoA4J7Cm8Q1e-NjkGEvCPMOqa4-KdCUjyGLnGTAPlDk,9957
 returnn/tf/__init__.py,sha256=X4g2LFCFTl0uiybMRkfBY8AYkgMa6HX0vVxxTk0nMiE,88
 returnn/tf/compat.py,sha256=NkAkdlR37m2d9qh3i33sIfEGilOaFBeCofAQpQwnZpY,1632
 returnn/tf/data_pipeline.py,sha256=iNkNHv5PiGcudlajG8eO336rPD3hya5kWMDrjhWa4jA,36632
@@ -172,13 +172,13 @@ returnn/tf/distributed.py,sha256=PCLspuNg4XP4ZX3Q444IlohUJEy0Dc8rp8YlmDqVbEc,151
 returnn/tf/engine.py,sha256=nhAMSEVUIf6Onm8jaRkT2CuY5XbOV5CEeWyOMUk67kY,146610
 returnn/tf/horovod.py,sha256=Dpv_3wZxB8q8Gqk6xah4iJ4vKGKWWg1-7PPhpSMPlec,5404
 returnn/tf/hyper_param_tuning.py,sha256=IfVRYYz-oSwOa2E7-vwh-pnWL4j-StHHbSYv7VbvcPE,31619
-returnn/tf/native_op.py,sha256=jieVDsDzF3cPsxncfy7lMheWf076FRbAE3L0P3TjwiY,79691
-returnn/tf/network.py,sha256=9kB82cTpdTtRW0mObYE7rAb_YRRwSfPYkkUZKRuXxio,224656
+returnn/tf/native_op.py,sha256=LuwPj-0lMBDZeP9q79cI1FuIYudnMavUlKsTQWMrkl4,77382
+returnn/tf/network.py,sha256=ZBo5qXOZHBJLjv2E8y9APeiRIpz5KEQdc6GN3rl6LBM,224668
 returnn/tf/sprint.py,sha256=Yqjh0-6sCWHpdDPQCzHKx7TwQCOjJyjfd0KHtnYdd-8,5471
 returnn/tf/updater.py,sha256=RcvoGnjBcObbLfLHH_mDRSY2lTeLyNoAFsZpHUiIgRY,72036
 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=ZHfmVD8uN0yeyqXFPxYxpnfBbjIpGQl9ykKKW9hxdD0,47537
+returnn/tf/frontend_layers/_backend.py,sha256=_YYQ-lV2srx-DwKIiaKacYkINOisPIstaPzzC7eRsaY,47652
 returnn/tf/frontend_layers/_utils.py,sha256=ijByaDOqPDod5mZC9EoTkt8PHBEODXHsWbkwDOF9XW4,4205
 returnn/tf/frontend_layers/cond.py,sha256=bGd_g2tzpKXO218Xk-so59vFPJF-jF_ZvoZIU-1qBzw,14832
 returnn/tf/frontend_layers/config_entry_points.py,sha256=t01RWOiaZohzuqPXX-MLV0P5yCOfE0dz-9dZ77_pK4c,5751
@@ -194,13 +194,13 @@ returnn/tf/frontend_low_level/__init__.py,sha256=34469k3KzMUIGowxReOZnbf6WdTjxY7
 returnn/tf/frontend_low_level/_backend.py,sha256=Hv838I2eyOP2qVNWs5DJxseyxUbAET2lm0ZZcbW_CsE,24991
 returnn/tf/layers/__init__.py,sha256=Ngu-X84nWFgz7ndDu88DqoZ-5lUMMTQWH4g7N8pSoCg,72
 returnn/tf/layers/base.py,sha256=sUxEfh6WxaHWHG7O3cfxB6gG6YpEHkFKUJVayKvTBSI,152968
-returnn/tf/layers/basic.py,sha256=PMYNoMq8qH41QhWhJPg5Uc409GZHkcnecouorg9sqJY,615466
+returnn/tf/layers/basic.py,sha256=jKzfRhBJgt5_tgIATde2kdza5u3aCqACx7BFyClngno,614277
 returnn/tf/layers/rec.py,sha256=3f6M_5aAMPvx7aAHdPV3VSFRHf7tjpp8lrXSzmk1I5c,548435
 returnn/tf/layers/segmental_model.py,sha256=wUyDZGr-eTVIIQWcsHLML0wtOxuWn_NFKOIrUKQcvoI,21515
 returnn/tf/layers/signal_processing.py,sha256=vRlkN7k7otk9_Qdv0qr_l6V0VT5Q6dO2MxwZWb2HH2M,52693
 returnn/tf/layers/variable.py,sha256=G1dIEoq0iQsXp-uOAUPTaBKHSOQfx7Sn-spD8MRv0HM,11446
 returnn/tf/util/__init__.py,sha256=mEg5jNVbQBLO2TGwO4Ff2F5qQN5_Zg4hAAQfX5taeec,92
-returnn/tf/util/basic.py,sha256=HWoRCkiPsMY0z9l0boB1NW0xqubVw-W8cZTZFHphNoc,303027
+returnn/tf/util/basic.py,sha256=ezK-XBKQcscVKmCL43wsieiUZntHUVOZwhDZtVItOqg,304088
 returnn/tf/util/data.py,sha256=AlSa0r_IaXtjKG1q1vxUybFazpjt4lUX8LYq0STJv-w,29471
 returnn/tf/util/gradient_checkpoint.py,sha256=_1NGAmNZ5NiGhFYVRWvBV5yejt-EZWbbvxNWHbESp5Q,7426
 returnn/tf/util/ken_lm.py,sha256=R60UAoywriuDIeQ2Hk3Vm_waf2Hxxc88ofzEw6X6Sd4,17313
@@ -208,7 +208,7 @@ returnn/tf/util/open_fst.py,sha256=sZRDw4TbxvhGqpGdUJWy1ebvlZm4_RPhygpRw9uLAOQ,1
 returnn/torch/README.md,sha256=jzJ2FpOHW02vxN69yKaV97C9LI-hmvjBglKfdZXIDdc,85
 returnn/torch/__init__.py,sha256=MHEUyNHB20Vy89uKAqZoj6FxJKF1Gq3HW-i6ra1pNcI,24
 returnn/torch/distributed.py,sha256=_lyJR71HIoCHpMi5GztGM7YwrX54Am8zSkjnDkE1Lbk,7524
-returnn/torch/engine.py,sha256=Lj_Go3Q5oHOVrHdNcz_6YLjODsiGUsJVMbU0EzSDByk,79677
+returnn/torch/engine.py,sha256=XaJhVpF181sf8M1iXAs3u0zr37VVUG3SW81-DIZgg3g,81280
 returnn/torch/updater.py,sha256=nNd1mBPQyvIB096BEFi0KKmRI-U3jnRETzb743p2B9c,32064
 returnn/torch/data/__init__.py,sha256=6cLNEi8KoGI12PF6akN7mI_mtjlx-0hcQAfMYoExwik,132
 returnn/torch/data/extern_data.py,sha256=5al706ZaYtHWLp5VH2vS-rW69YXP3NHyOFRKY0WY714,7810
@@ -217,9 +217,9 @@ returnn/torch/data/queued_data_iter.py,sha256=PoOsGHdHVZjTmcyfq_ZOw--P6hyfTdmAWI
 returnn/torch/data/returnn_dataset_wrapper.py,sha256=fMahf05G0SPYm6HxSQpVm8JhsIHons-i1Ce4aQv4IjM,8332
 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=2otTWQwS-PGRJImGqUAikChP3JcOv_KHgPAQX4VycJg,103382
+returnn/torch/frontend/_backend.py,sha256=8EBRGN0jY5rl9Z5-wd4kvoDesssWcVDVXNl25-bG8cA,108882
 returnn/torch/frontend/_rand.py,sha256=1JgIkV2XmpgJD86zXZ-NCAe-QuoP2swr6NaS1oz3Qa8,1830
-returnn/torch/frontend/bridge.py,sha256=c_mVBCBo29sjm8Bhxarv00szwGPgxjwoIqAHOmceGQw,7842
+returnn/torch/frontend/bridge.py,sha256=RBtAIlYWn_AC-GaHWperrOncPjMLWAOrU30pWk2789A,9775
 returnn/torch/frontend/raw_ops.py,sha256=lF0h-KtYYsdaaqQADylVZp9qzPskOOXA4MfmYDyx5IU,296
 returnn/torch/optim/README.md,sha256=0iH5FiKb7iDrVK5n8V6yCh4ciCFG2YSbyh7lPneT5ik,360
 returnn/torch/optim/__init__.py,sha256=yxdbnOkXAHzZ_t6cHi6zn5x_DQNlLZJ-KxZByHTIg1U,29
@@ -229,18 +229,18 @@ returnn/torch/util/__init__.py,sha256=AOXYUjzPm0XrzFJCPAXo9Jj_FvqD1XH3FfKtho80Vl
 returnn/torch/util/array_.py,sha256=ell3VZvn01SLtF9Pw2fvPzFNO-XDQ7tSB9VCrVSKmSA,2556
 returnn/torch/util/debug_inf_nan.py,sha256=fmzSSTJJyLf7i5yDWRHLeDI0gxvadeqLE8RxMuSHx_4,6398
 returnn/torch/util/diagnose_gpu.py,sha256=_yswLmwR8Q2rCsv2jI5FUQNBT__453jBmiWYwazdu20,6808
-returnn/torch/util/exception_helper.py,sha256=_SqxTD5F-GDY2eR4uRALyUTJwt0ytcbJGB_w38RJMBA,4320
+returnn/torch/util/exception_helper.py,sha256=54IzlsXYp6E_rEEWIpgppkFid9stb-2PZVRU8d5mFNE,4497
 returnn/torch/util/gradient_checkpoint.py,sha256=iLy-FB65DC8O6LxzmMvFjnSdpIVpko87ppIvRKAbtpQ,27995
 returnn/torch/util/module.py,sha256=MXHIrF9Isu575DDJIa81212ULKwdqu1oOLxDVZecVSk,1693
 returnn/torch/util/scaled_gradient.py,sha256=C5e79mpqtxdtw08OTSy413TSBSlOertRisc-ioiFIaU,3191
 returnn/util/__init__.py,sha256=UIG1qw4idqhW71BV60ha7h9PktxvEVcBIu0lYRossK8,336
-returnn/util/basic.py,sha256=MBbyYddO0aJj8Xcb9RbQ5VeunC7htc9IGex09eao67o,143323
-returnn/util/better_exchook.py,sha256=39yvRecluDgYhViwSkaQ8crJ_cBWI63KeEGuK4RKe5w,70843
+returnn/util/basic.py,sha256=rFeg3XwjNcNDbBgjkhisStbjTFA8CEfIrdwHjfdkJKw,143212
+returnn/util/better_exchook.py,sha256=hOKazwv2q2-d0XMfxkJXMbLZyNTtraV3jPHplFcrMsg,71014
 returnn/util/bpe.py,sha256=LWFhICZsEOnMwNws0lybPNzKRX6rSr8yKCvP65vjl9Y,19656
 returnn/util/collect_outputs_dict.py,sha256=CjpsftoMgmvyE4wNKTO6F-QQ_44QHXcOZIXMUMQVZ-8,2637
-returnn/util/debug.py,sha256=wuRzdg9zB84WWCGyTjmRR_zYypu8gXxlc0nZ6si9OC8,28224
+returnn/util/debug.py,sha256=0ED4etMKG9lVqU0HPKEiCK-HoS8hBgnQza444QCE6ec,28576
 returnn/util/debug_helpers.py,sha256=0EINLK4uLtoSt5_kHs1M2NIFpMd0S7i4c4rx90U4fJk,2914
-returnn/util/file_cache.py,sha256=Wyx9u1P59PAo4HouVPh6uaGQeFIXGZHshDyvcRn4YlM,28786
+returnn/util/file_cache.py,sha256=8xE4zMQi38g7ZIGwNohd13_CgjzpIs18ILxFCKttzxE,29439
 returnn/util/fsa.py,sha256=k2lJ8tyf_g44Xk1EPVLwDwpP4spoMTqIigDVOWocQHY,59177
 returnn/util/literal_py_to_pickle.py,sha256=3dnjWPeeiDT2xp4bRDgIf9yddx7b1AG7mOKEn_jiSl8,2173
 returnn/util/lru_cache.py,sha256=7Q5H3a8b07E8e1iB7PA9jCpRnxMJZOFS2KO07cy0gqk,11446
@@ -252,11 +252,11 @@ returnn/util/py-to-pickle.cpp,sha256=ByU4cwy5MGEihaoYiRo1sSsJfYn10_riDwVqSHRLwp8
 returnn/util/py_ext_mod_compiler.py,sha256=I1w9laIPqJbQGb2lFp-3llBjORS-217ZGIbPCp6PIes,1708
 returnn/util/result_with_reason.py,sha256=6jS7caYrZADrb8o-CpQnTfskZb3fTNMcKU-JlnIh6Kg,359
 returnn/util/sig_proc.py,sha256=Tjz0VOAVyqu2qDCF5HZ1JjALjcFsHcNkcd96WgZeKfE,7265
-returnn/util/task_system.py,sha256=y4sMVXQ25Qd2z0rx03uOlXlkE-jbCYC1Sjfn-XlraVU,26003
+returnn/util/task_system.py,sha256=7Dz7Nvi_1-o5pDv9OZYdAnlJw6OSvgbYUmQ72P0Fgkw,26002
 returnn/util/train_proc_manager.py,sha256=Pjht28k6uz6BNQ47uW6Gf880iyq5q4wx7P_K2tmoAM8,3266
 returnn/util/watch_memory.py,sha256=BR5P2kvBN6UI81cE0_1WAA6Hd1SByLbBaiDxvLhPOew,4213
-returnn-1.20251027.224345.dist-info/LICENSE,sha256=ywBD_U2aD4vpuoIgNAsjIGBYydl0tVKll3De0Z8s77c,11041
-returnn-1.20251027.224345.dist-info/METADATA,sha256=GJaNdcwgbvQQFiUHTgzu_yvcMXGo3HTUjilKYreP16c,5215
-returnn-1.20251027.224345.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
-returnn-1.20251027.224345.dist-info/top_level.txt,sha256=Lsn4WZc5Pbfk0-xDQOgnFCxOoqxL4CyeM3N1TFbJncw,8
-returnn-1.20251027.224345.dist-info/RECORD,,
+returnn-1.20260109.93428.dist-info/LICENSE,sha256=ywBD_U2aD4vpuoIgNAsjIGBYydl0tVKll3De0Z8s77c,11041
+returnn-1.20260109.93428.dist-info/METADATA,sha256=8G2OFR-V5IlE98f0vmLneA27jg9-B7eN973G7vJpj0I,5215
+returnn-1.20260109.93428.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
+returnn-1.20260109.93428.dist-info/top_level.txt,sha256=Lsn4WZc5Pbfk0-xDQOgnFCxOoqxL4CyeM3N1TFbJncw,8
+returnn-1.20260109.93428.dist-info/RECORD,,