returnn 1.20250508.93313__py3-none-any.whl → 1.20250513.145447__py3-none-any.whl

returnn/frontend/_utils.py

@@ -263,7 +263,7 @@ def _slice_find_sparse_dim(v: Union[Tensor, slice, Any]) -> Optional[Dim]:
 
 
 def _map_slice_value_raw(
-    v: Union[None, slice, int, numpy.number, numpy.ndarray, Tensor[T]]
+    v: Union[None, slice, int, numpy.number, numpy.ndarray, Tensor[T]],
 ) -> Union[None, slice, int, numpy.number, T]:
     if v is None:
         return None

returnn/frontend/audio/mel.py

@@ -2,7 +2,6 @@
 Mel filterbank.
 """
 
-
 from __future__ import annotations
 from typing import Optional, Union, Tuple
 import functools

returnn/frontend/const.py

@@ -54,9 +54,9 @@ def full(
             "Use rf.convert_to_tensor to convert an arbitrary array to a tensor."
         )
     if isinstance(fill_value, Tensor):
-        assert (
-            fill_value.dims == ()
-        ), f"full/fill/constant: expect scalar fill_value, got tensor with shape {fill_value.dims}."
+        assert fill_value.dims == (), (
+            f"full/fill/constant: expect scalar fill_value, got tensor with shape {fill_value.dims}."
+        )
     return global_backend.full(
         dims, fill_value, dtype=dtype, device=device, sparse_dim=sparse_dim, feature_dim=feature_dim
     )

returnn/frontend/device.py

@@ -2,7 +2,6 @@
 Device handling.
 """
 
-
 from __future__ import annotations
 from typing import Optional
 from contextlib import contextmanager

returnn/frontend/dropout.py

@@ -50,7 +50,7 @@ def dropout(
         raise ValueError(f"dropout axis {axis} not in source {source}")
 
     if isinstance(keep_prob, (float, int)) and not 0 < keep_prob <= 1:
-        raise ValueError("keep_prob must be a scalar tensor or a float in the " "range (0, 1], got %g" % keep_prob)
+        raise ValueError("keep_prob must be a scalar tensor or a float in the range (0, 1], got %g" % keep_prob)
 
     # Do nothing if we know keep_prob == 1
     if isinstance(keep_prob, (float, int)) and keep_prob == 1:

returnn/frontend/encoder/e_branchformer.py

@@ -268,7 +268,7 @@ class Merge(rf.Module):
 
 
 def _make_activation(
-    activation: Union[Callable[[Tensor], Tensor], Dict[str, Any], rf.Module]
+    activation: Union[Callable[[Tensor], Tensor], Dict[str, Any], rf.Module],
 ) -> Union[Callable[[Tensor], Tensor], rf.Module]:
     if isinstance(activation, dict):
         activation = rf.build_from_dict(activation)

returnn/frontend/loop.py

@@ -273,9 +273,9 @@ def _check_matching_loop_var_templates(loop_var_templates: S, loop_vars: S):
             x._push_back_delayed_check()
 
         else:  # other cases: just check same type
-            assert type(template) is type(
-                x
-            ), f"loop var {path} template type {type(template)} does not match var type {type(x)}"
+            assert type(template) is type(x), (
+                f"loop var {path} template type {type(template)} does not match var type {type(x)}"
+            )
             assert not isinstance(x, Tensor), f"loop var {path} is a Tensor but should not be"
 
     tree.map_structure_with_path(_check, loop_var_templates, loop_vars)

returnn/frontend/loss.py

@@ -137,7 +137,6 @@ def edit_distance(a: Tensor, a_spatial_dim: Dim, b: Tensor, b_spatial_dim: Dim,
     # We are going diagonal over (Ta+1) and (Tb+1). (Similar as RETURNN native EditDistanceOp.)
     # You need to draw the grid on paper to understand all the index math...
     for u in range(1, n_a_max_len + n_b_max_len + 1):
-
         prev2_dist, _ = rf.slice(
             buffer, axis=buffer_dim, start=buffer_offsets[u % 3], size=b_spatial_dim1, out_dim=b_spatial_dim1
         )  # [Tb+1,B]

returnn/frontend/matmul.py

@@ -2,7 +2,6 @@
 Dot / matmul
 """
 
-
 from __future__ import annotations
 from typing import Sequence, Union, TypeVar
 from returnn.tensor import Tensor, Dim

returnn/frontend/run_ctx.py

@@ -306,19 +306,19 @@ class RunCtx:
         assert self.stage == "forward_step"
 
         if self.expected_outputs is not None:
-            assert (
-                name in self.expected_outputs.data
-            ), f"mark_as_output: unexpected output {name!r}, we expect outputs: {self.expected_outputs}"
+            assert name in self.expected_outputs.data, (
+                f"mark_as_output: unexpected output {name!r}, we expect outputs: {self.expected_outputs}"
+            )
         expected_output = self.expected_outputs.data[name] if self.expected_outputs else None
-        assert dims is None or (
-            isinstance(dims, (list, tuple)) and all(isinstance(dim, Dim) for dim in dims)
-        ), f"dims should be a tuple of Dims, got {dims}"
+        assert dims is None or (isinstance(dims, (list, tuple)) and all(isinstance(dim, Dim) for dim in dims)), (
+            f"dims should be a tuple of Dims, got {dims}"
+        )
         if dims is None and expected_output is not None:
             dims = expected_output.dims
         if dims is not None and expected_output is not None:
-            assert expected_output.dims == tuple(
-                dims
-            ), f"mark_as_output: {name!r} dims mismatch from expected output, given {dims}, expected {expected_output}"
+            assert expected_output.dims == tuple(dims), (
+                f"mark_as_output: {name!r} dims mismatch from expected output, given {dims}, expected {expected_output}"
+            )
 
         if not isinstance(tensor, Tensor):
             assert isinstance(tensor, _backend.global_backend.RawTensorType)

returnn/frontend/signal.py

@@ -2,7 +2,6 @@
 stft etc
 """
 
-
 from __future__ import annotations
 from typing import Optional, Tuple
 from returnn.tensor import Tensor, Dim

returnn/frontend/types.py

@@ -19,15 +19,13 @@ ItemKeyType = Union[RawTensorTypes, Tensor, slice, Sequence[Union[RawTensorTypes
 class GetModelFunc(Protocol):
     """get model func"""
 
-    def __call__(self, *, epoch: int, step: int) -> rf.Module:
-        ...
+    def __call__(self, *, epoch: int, step: int) -> rf.Module: ...
 
 
 class StepFunc(Protocol):
     """step func"""
 
-    def __call__(self, *, model: rf.Module, extern_data: TensorDict) -> None:
-        ...
+    def __call__(self, *, model: rf.Module, extern_data: TensorDict) -> None: ...
 
 
 def get_raw_tensor_type() -> Type:

returnn/native_op.py

@@ -291,6 +291,7 @@ class LstmGenericBase(NativeOpGenBase):
       :param H: gates and cell state. 3d (time,batch,dim*4)
       :param d: final cell state. 2d (batch,dim)
     """
+
     in_info = (
         {
             "name": "Z",
@@ -542,6 +543,7 @@ class LstmLowMem(NativeOpGenBase):
       :param C: cell states. 3d (time,batch,dim). gradient ignored!
       :param d: final cell state. 2d (batch,dim)
     """
+
     in_info = (
         {"name": "X", "ndim": 3, "shape": (None, None, None), "need_contiguous": True},
         {"name": "W", "ndim": 2, "shape": (None, None), "need_contiguous": True},
@@ -994,6 +996,7 @@ class NativeLstm2(NativeOpGenBase):
       :param H: cell-in + gates. 3d (time,batch,dim*4). gradient ignored!
       :param d: final cell state. 2d (batch,dim)
     """
+
     in_info = (
         {"name": "X", "ndim": 3, "shape": (None, None, None), "need_contiguous": True},
         {"name": "W", "ndim": 2, "shape": (None, None), "need_contiguous": True},
@@ -1423,6 +1426,7 @@ class TwoDLSTM(NativeOpGenBase):
       :param H: gates and cell state. 4d (timeS,timeT,batch,dim*5) ?
       :param d: final cell state. 3d (timeT,batch,dim)
     """
+
     in_info = (
         {
             "name": "X",
@@ -3198,6 +3202,7 @@ class FastBaumWelchOp(NativeOpGenBase):
     outputs:
       :param output: Baum-Welch alignment, scores in -log space. 3d (time,batch,dim), like am_scores
     """
+
     in_info = (
         {
             "name": "am_scores",
@@ -3620,6 +3625,7 @@ class MultiEndFastBaumWelchOp(NativeOpGenBase):
     outputs:
       :param output: Baum-Welch alignment, scores in -log space. 3d (time,batch,dim), like am_scores
     """
+
     in_info = (
         {
             "name": "am_scores",
@@ -4497,6 +4503,7 @@ class FastViterbiOp(NativeOpGenBase):
       :param output: Viterbi (hard) alignment, scores in +log space. 2d (time,batch)
       :param scores: (batch,)
     """
+
     in_info = (
         {
             "name": "am_scores",
@@ -4865,6 +4872,7 @@ class GetCtcFsaFastBwOp(NativeOpGenBase):
     `num_edges` should be `n_batch * (5 * (n_time - 1) + 10)`
       (see construction in kernel why that number).
     """
+
     in_info = (
         {
             "name": "targets",
@@ -5229,6 +5237,7 @@ class EditDistanceOp(NativeOpGenBase):
     outputs:
       :param output: 1d (batch,), int32, unnormalized edit distance
     """
+
     in_info = (
         {
             "name": "a",
@@ -5414,6 +5423,7 @@ class OptimalCompletionEditDistanceOp(NativeOpGenBase):
     outputs:
       :param output: 1d (batch,), int32, unnormalized edit distance
     """
+
     in_info = (
         {
             "name": "a",
@@ -5610,6 +5620,7 @@ class OptimalCompletionEditDistancePerSuccessorOp(NativeOpGenBase):
     outputs:
       :param output: 2d (batch,num_labels), int32, unnormalized edit distance
     """
+
     in_info = (
         {
             "name": "a",
@@ -5880,6 +5891,7 @@ class NextEditDistanceRowOp(NativeOpGenBase):
     outputs:
       :param output: 2d (batch,b_time + 1), int32, next (unnormalized) edit distance row
     """
+
     in_info = (
         {
             "name": "last_row",
@@ -6039,6 +6051,7 @@ class NextEditDistanceReduceOp(NativeOpGenBase):
     outputs:
       :param output: 2d (batch,n_labels), int32, next (unnormalized) (maybe optional) edit distance
     """
+
     in_info = (
         {
             "name": "last_row",

returnn/sprint/cache.py

@@ -7,7 +7,7 @@ This module is about reading (maybe later also writing) the Sprint archive forma
 """
 
 from __future__ import annotations
-from typing import List
+from typing import List, Optional, Tuple
 import sys
 import os
 import typing
@@ -904,9 +904,7 @@ class MixtureSet:
             self.densities[n, 1] = cov_idx
 
         self.num_mixtures = self.read_u32()
-        self.mixtures = [
-            None
-        ] * self.num_mixtures  # type: typing.List[typing.Optional[typing.Tuple[typing.List[int],typing.List[float]]]]  # nopep8
+        self.mixtures: List[Optional[Tuple[List[int], List[float]]]] = [None] * self.num_mixtures
         for n in range(self.num_mixtures):
             num_densities = self.read_u32()
             dns_idx = []

returnn/sprint/interface.py

@@ -820,9 +820,9 @@ def _prepare_forwarding():
     assert engine
     assert config
     # Should already be set via setTargetMode().
-    assert config.list("extract") == [
-        "posteriors"
-    ], "You need to have extract = posteriors in your RETURNN config. You have: %s" % config.list("extract")
+    assert config.list("extract") == ["posteriors"], (
+        "You need to have extract = posteriors in your RETURNN config. You have: %s" % config.list("extract")
+    )
 
     # Load network.
     engine.init_network_from_config(config)
@@ -870,7 +870,6 @@ def _train(segment_name, features, targets=None):
     # The CRNN train thread started via start() will do the actual training.
 
     if TargetMode == "criterion-by-sprint":
-
         # TODO...
         make_criterion_class()
 

returnn/tensor/_dim_extra.py

@@ -1067,13 +1067,14 @@ class _DimMixin:
                     )
                 )
         if batch and getattr(x, "_RETURNN_dyn_size_beam", None):
-            assert batch.beam == getattr(
-                x, "_RETURNN_dyn_size_beam"
-            ), "%s: dyn size %s has unexpected batch %s, expected %s" % (
-                self,
-                x,
-                batch,
-                getattr(x, "_RETURNN_dyn_size_beam"),
+            assert batch.beam == getattr(x, "_RETURNN_dyn_size_beam"), (
+                "%s: dyn size %s has unexpected batch %s, expected %s"
+                % (
+                    self,
+                    x,
+                    batch,
+                    getattr(x, "_RETURNN_dyn_size_beam"),
+                )
             )
         if self.batch and batch:
             assert self.batch == batch
@@ -1359,8 +1360,7 @@ class _DimMixin:
                 # Only auto-generated dim tags are allowed to be treated as broadcastable.
                 # This was another suggestion from here: https://github.com/rwth-i6/returnn/issues/666
                 # It was not implemented like this because the auto_generated flag was only introduced later.
-                (self.dimension == 1 and self.auto_generated)
-                or (other.dimension == 1 and other.auto_generated)
+                (self.dimension == 1 and self.auto_generated) or (other.dimension == 1 and other.auto_generated)
             ):
                 pass  # pass on
             else:

returnn/tensor/_tensor_extra.py

@@ -335,9 +335,9 @@ class _TensorMixin(_TensorMixinBase):
                         if tag.dyn_size_ext.placeholder is None:
                             tag.complete_dyn_size()
                         if self.placeholder is not None:
-                            assert (
-                                tag.dyn_size_ext.placeholder is not None
-                            ), "%s sanity_check: dynamic dim %s value unknown" % (self, tag)
+                            assert tag.dyn_size_ext.placeholder is not None, (
+                                "%s sanity_check: dynamic dim %s value unknown" % (self, tag)
+                            )
                 assert tag.is_dim_known()
 
     def get_runtime_sanity_check_op(self: Tensor):
@@ -2494,8 +2494,7 @@ class _TensorMixin(_TensorMixinBase):
                     if res_tag.match_priority > tag.match_priority:
                         continue
                     raise Exception(
-                        f"{self}: get_axis_from_description({axis}) not unique."
-                        f" use match_priority to resolve ambiguity"
+                        f"{self}: get_axis_from_description({axis}) not unique. use match_priority to resolve ambiguity"
                     )
             if res_idx is None:
                 raise Exception(f"{self}: get_axis_from_description({axis}) not found")
@@ -2646,12 +2645,13 @@ class _TensorMixin(_TensorMixinBase):
             return self.batch_shape[self.time_dim_axis_excluding_batch] is None
         if self.time_dim_axis_excluding_batch in self.size_placeholder:
             return True
-        assert isinstance(
-            self.shape[self.time_dim_axis_excluding_batch], int
-        ), "%s: dynamic time axis dim (None) (axis %i) but size_placeholder %r misses information" % (
-            self,
-            self.time_dim_axis,
-            self.size_placeholder,
+        assert isinstance(self.shape[self.time_dim_axis_excluding_batch], int), (
+            "%s: dynamic time axis dim (None) (axis %i) but size_placeholder %r misses information"
+            % (
+                self,
+                self.time_dim_axis,
+                self.size_placeholder,
+            )
         )
         return False
 
@@ -3307,14 +3307,15 @@ class _TensorMixin(_TensorMixinBase):
                     if self_axis not in taken_self_axes
                 ]
                 if opt == "unknown_spatial_matches":
-                    assert (
-                        len(matching) <= 1
-                    ), "cannot match axes %s from %s to %s, failed at other %s, not unique after %s" % (
-                        other_axes,
-                        other,
-                        self,
-                        other_axis,
-                        opt,
+                    assert len(matching) <= 1, (
+                        "cannot match axes %s from %s to %s, failed at other %s, not unique after %s"
+                        % (
+                            other_axes,
+                            other,
+                            self,
+                            other_axis,
+                            opt,
+                        )
                     )
                 if matching:
                     break

returnn/tensor/_tensor_op_overloads.py

@@ -13,7 +13,6 @@ from ._tensor_mixin_base import _TensorMixinBase
 
 
 class _TensorOpOverloadsMixin(_TensorMixinBase):
-
     # Note that all those ops have native implementations as well,
     # so keep the logic in sync.
 

returnn/tensor/tensor.py

@@ -187,7 +187,7 @@ class Tensor(_TensorMixin, _TensorOpOverloadsMixin, Generic[RawTensorType]):
             if not backend.executing_eagerly():
                 backend.set_known_shape_raw(value, self.batch_shape)
             assert backend.get_dtype_name_raw(value) == self.dtype, (
-                f"{self} dtype {self.dtype} does not match " f"raw tensor dtype {backend.get_dtype_name_raw(value)}"
+                f"{self} dtype {self.dtype} does not match raw tensor dtype {backend.get_dtype_name_raw(value)}"
             )
         self._raw_tensor = value
 

returnn/tensor/tensor_dict.py

@@ -91,9 +91,9 @@ class TensorDict:
         out = {}
         for key, value in self.data.items():
             assert key not in out
-            assert isinstance(
-                value.raw_tensor, expected_value_type
-            ), f"key {key} {value}: unexpected {type(value.raw_tensor)}, expected {expected_value_type}"
+            assert isinstance(value.raw_tensor, expected_value_type), (
+                f"key {key} {value}: unexpected {type(value.raw_tensor)}, expected {expected_value_type}"
+            )
             out[key] = value.raw_tensor
             for i, dim in enumerate(value.dims):
                 if exclude_duplicate_dims and dim in visited_dims:
@@ -103,9 +103,9 @@ class TensorDict:
                 if dim.is_batch_dim() and (dim.dyn_size_ext is None or dim.dyn_size_ext.raw_tensor is None):
                     if include_scalar_dyn_sizes:
                         dim_value = dim.get_dim_value()
-                        assert isinstance(
-                            dim_value, expected_value_type
-                        ), f"key {key_} {dim}: unexpected {type(dim_value)}, expected {expected_value_type}"
+                        assert isinstance(dim_value, expected_value_type), (
+                            f"key {key_} {dim}: unexpected {type(dim_value)}, expected {expected_value_type}"
+                        )
                         out[key_] = dim_value
                 elif dim.dyn_size_ext is not None:
                     if include_scalar_dyn_sizes or dim.dyn_size_ext.dims:
@@ -116,9 +116,9 @@ class TensorDict:
                         out[key_] = dim.dyn_size_ext.raw_tensor
                 elif dim.size is not None:
                     if include_scalar_dyn_sizes and include_const_sizes:
-                        assert isinstance(
-                            dim.size, expected_value_type
-                        ), f"key {key_} {dim}: unexpected {type(dim.size)}, expected {expected_value_type}"
+                        assert isinstance(dim.size, expected_value_type), (
+                            f"key {key_} {dim}: unexpected {type(dim.size)}, expected {expected_value_type}"
+                        )
                         out[key_] = dim.size
                 else:
                     raise Exception(f"cannot handle dim: {dim}")