returnn 1.20250226.115259__tar.gz → 1.20250226.183415__tar.gz

{returnn-1.20250226.115259/returnn.egg-info → returnn-1.20250226.183415}/PKG-INFO

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

returnn-1.20250226.183415/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20250226.183415'
+long_version = '1.20250226.183415+git.ba9d72e'

{returnn-1.20250226.115259 → returnn-1.20250226.183415}/returnn/frontend/_numpy_backend.py

@@ -159,6 +159,35 @@ class NumpyBackend(Backend[numpy.ndarray]):
             res = numpy.array(res)
         return res
 
+    @staticmethod
+    def where(
+        cond: Tensor,
+        true_: Union[Tensor, rf.RawTensorTypes],
+        false_: Union[Tensor, rf.RawTensorTypes],
+        *,
+        allow_broadcast_all_sources: bool = False,
+    ) -> Tensor:
+        """where"""
+        if isinstance(true_, Tensor):
+            dtype = true_.dtype
+        elif isinstance(false_, Tensor):
+            dtype = false_.dtype
+        else:
+            dtype = None
+        true_ = rf.convert_to_tensor(true_, _backend=NumpyBackend, dtype=dtype)
+        false_ = rf.convert_to_tensor(false_, _backend=NumpyBackend, dtype=dtype)
+        out = Tensor.get_common_data(
+            [true_, false_, cond], allow_broadcast_all_sources=allow_broadcast_all_sources, name="where"
+        )
+        out.dtype = true_.dtype
+        out.sparse_dim = true_.sparse_dim or false_.sparse_dim
+        out.feature_dim = true_.feature_dim or false_.feature_dim
+        cond_bc_raw = cond.copy_compatible_to_dims_raw(out.dims)
+        true_bc_raw = true_.copy_compatible_to_dims_raw(out.dims)
+        false_bc_raw = false_.copy_compatible_to_dims_raw(out.dims)
+        out.raw_tensor = numpy.where(cond_bc_raw, true_bc_raw, false_bc_raw)
+        return out
+
     @staticmethod
     def range_over_dim(dim: Dim, *, dtype: Optional[str] = None, device: Optional[str] = None) -> Tensor[numpy.ndarray]:
         """

{returnn-1.20250226.115259 → returnn-1.20250226.183415}/returnn/frontend/array_.py

@@ -325,6 +325,7 @@ def window(
     padding: str = "same",
     pad_value: Optional[Union[int, float]] = None,
     stride: int = 1,
+    use_mask: Optional[bool] = None,
 ) -> Tuple[Tensor, Dim]:
     """
     Follows the same idea as RETURNN tf_util.windowed,
@@ -338,8 +339,14 @@ def window(
     :param padding: "same" or "valid"
     :param pad_value:
     :param stride:
+    :param use_mask: whether we should mask to make sure the zero padding is correct
     :return: out, out_spatial_dim
     """
+    if spatial_dim.need_masking():
+        if use_mask is None:
+            use_mask = rf.use_mask_default(default=True, default_false_for_behavior_version_up_to=22)
+        if use_mask:
+            source = source.copy_masked(0, dims=[spatial_dim])
     assert window_dim.dimension is not None
     if padding == "same":
         out_spatial_dim = spatial_dim
@@ -784,6 +791,7 @@ def scatter(
     mode: str = "sum",
     fill_value: Optional[Union[int, float]] = None,
     out_dim: Optional[Union[Dim, Sequence[Dim]]] = None,
+    use_mask: Optional[bool] = None,
 ) -> Tensor:
     """
     Scatters into new zero-tensor.
@@ -806,6 +814,7 @@ def scatter(
         If not given, will be automatically determined as the sparse_dim from indices.
         If multiple out dims, use indices into the merged out dims,
         and then we use :func:`rf.split_dims` afterwards.
+    :param use_mask:
     :return: [batch_dims..., out_dim(s)..., feature_dims...]
     """
     if mode == "logsumexp":
@@ -826,24 +835,32 @@ def scatter(
             fill_value = 0
         elif mode == "max":
             if "int" in source.dtype:
-                import numpy
-
                 fill_value = numpy.iinfo(source.raw_tensor.dtype).min
             else:
                 fill_value = float("-inf")
         elif mode == "min":
             if "int" in source.dtype:
-                import numpy
-
                 fill_value = numpy.iinfo(source.raw_tensor.dtype).max
             else:
                 fill_value = float("inf")
         else:
             raise ValueError(f"scatter: invalid mode {mode!r}")
+    indices_dim = indices_dim if isinstance(indices_dim, (list, tuple)) else [indices_dim]
+    if any(dim.need_masking() for dim in indices_dim):
+        if use_mask is None:
+            use_mask = rf.use_mask_default(default=True, default_false_for_behavior_version_up_to=22)
+        if use_mask:
+            source = source.copy_masked(fill_value, dims=indices_dim)
+    else:
+        use_mask = False
     # noinspection PyProtectedMember
-    return source._raw_backend.scatter(
+    out = source._raw_backend.scatter(
         source, indices=indices, indices_dim=indices_dim, mode=mode, fill_value=fill_value, out_dim=out_dim
     )
+    if use_mask and mode != "sum":
+        # Make sure we don't leave any infinities in the output.
+        out = out.copy_masked(0, dims=[out_dim])
+    return out
 
 
 def scatter_argmax(

{returnn-1.20250226.115259 → returnn-1.20250226.183415}/returnn/frontend/conv.py

@@ -34,8 +34,8 @@ class _ConvOrTransposedConv(rf.Module):
     Base class for both convolution and transposed convolution.
     """
 
-    nd: Optional[int] = None
-    _transposed: bool
+    nd: Optional[int] = None  # set in the subclasses, e.g. 1 for Conv1d, etc
+    _transposed: bool  # set in the subclasses _Conv or _TransposedConv
     groups: Optional[int] = None
 
     def __init__(
@@ -187,8 +187,14 @@ def conv(
     dilation_rate: Optional[Union[int, Sequence[int]]] = None,
     groups: Optional[int] = None,
     bias: Optional[Tensor] = None,
+    use_mask: Optional[bool] = None,
 ) -> Tuple[Tensor, Sequence[Dim]]:
     """convolution"""
+    if any(in_spatial_dim.need_masking() for in_spatial_dim in in_spatial_dims):
+        if use_mask is None:
+            use_mask = rf.use_mask_default(default=True, default_false_for_behavior_version_up_to=22)
+        if use_mask:
+            source = source.copy_masked(0, dims=in_spatial_dims)
     for in_spatial_dim in in_spatial_dims:
         if in_spatial_dim not in source.dims:
             raise ValueError(f"conv: source {source} does not have spatial dim {in_spatial_dim}")
@@ -345,8 +351,14 @@ def transposed_conv(
     output_padding: Optional[Union[Sequence[Optional[int]], int]] = None,
     strides: Optional[Sequence[int]] = None,
     bias: Optional[Tensor] = None,
+    use_mask: Optional[bool] = None,
 ) -> Tuple[Tensor, Sequence[Dim]]:
     """transposed conv"""
+    if any(in_spatial_dim.need_masking() for in_spatial_dim in in_spatial_dims):
+        if use_mask is None:
+            use_mask = rf.use_mask_default(default=True, default_false_for_behavior_version_up_to=22)
+        if use_mask:
+            source = source.copy_masked(0, dims=in_spatial_dims)
     # noinspection PyProtectedMember
     out, out_spatial_dims = source._raw_backend.transposed_conv(
         source=source,
@@ -394,6 +406,7 @@ class TransposedConv3d(_TransposedConv):
 def pool(
     source: Tensor,
     *,
+    nd: Optional[int] = None,
     mode: str,
     pool_size: Union[Sequence[int], int],
     padding: str = "valid",
@@ -401,22 +414,23 @@ def pool(
     strides: Optional[Union[Sequence[int], int]] = None,
     in_spatial_dims: Union[Sequence[Dim], Dim],
     out_spatial_dims: Optional[Union[Sequence[Dim], Dim]] = None,
-    nd: Optional[int] = None,
+    use_mask: Optional[bool] = None,
 ) -> Tuple[Tensor, Sequence[Dim]]:
     """
     A generic N-D pooling layer.
     This would usually be done after a convolution for down-sampling.
 
-    :param Tensor source:
+    :param source:
     :param nd:
-    :param str mode: "max" or "avg"
-    :param tuple[int] pool_size: shape of the window of each reduce
-    :param str padding: "valid" or "same"
-    :param tuple[int]|int dilation_rate:
-    :param tuple[int]|int|None strides: in contrast to tf.nn.pool, the default (if it is None) will be set to pool_size
-    :param Sequence[Dim] in_spatial_dims:
-    :param Sequence[Dim]|None out_spatial_dims:
-    :return: layer, out_spatial_dims
+    :param mode: "max" or "avg"
+    :param pool_size: shape of the window of each reduce
+    :param padding: "valid" or "same"
+    :param dilation_rate:
+    :param strides: in contrast to tf.nn.pool, the default (if it is None) will be set to pool_size
+    :param in_spatial_dims:
+    :param out_spatial_dims:
+    :param use_mask:
+    :return: out, out_spatial_dims
     """
     if isinstance(in_spatial_dims, Dim):
         in_spatial_dims = [in_spatial_dims]
@@ -440,6 +454,14 @@ def pool(
     assert isinstance(strides, (list, tuple))
     assert len(strides) == nd
 
+    if any(in_spatial_dim.need_masking() for in_spatial_dim in in_spatial_dims):
+        if use_mask is None:
+            use_mask = rf.use_mask_default(default=True, default_false_for_behavior_version_up_to=22)
+        if use_mask:
+            source = source.copy_masked({"max": float("-inf"), "avg": 0}[mode], dims=in_spatial_dims)
+    else:
+        use_mask = False
+
     # noinspection PyProtectedMember
     out, out_spatial_dims = source._raw_backend.pool(
         source=source,
@@ -451,6 +473,10 @@ def pool(
         in_spatial_dims=in_spatial_dims,
         out_spatial_dims=out_spatial_dims,
     )
+    if use_mask and mode == "max":
+        # We masked with -inf for max-pooling to get correct pooling at the boundaries.
+        # However, the resulting tensor might have -inf in it, and it is better to mask it out.
+        out = out.copy_masked(0, dims=out_spatial_dims)
     return out, out_spatial_dims
 
 

{returnn-1.20250226.115259 → returnn-1.20250226.183415}/returnn/frontend/dims.py

@@ -22,6 +22,7 @@ __all__ = [
     "num_elements_of_shape",
     "masked_fraction_of_shape",
     "last_frame_position_of_dim",
+    "use_mask_default",
 ]
 
 
@@ -305,3 +306,42 @@ def last_frame_position_of_dim(
     pos = rf.maximum(pos, 0)
     pos.sparse_dim = dim
     return pos
+
+
+def use_mask_default(
+    *, default: Optional[bool] = None, default_false_for_behavior_version_up_to: Optional[int] = None
+) -> Optional[bool]:
+    """
+    Check the global RETURNN config for the ``rf_use_mask``
+    on what default we should use for the ``use_mask`` argument in various functions
+    (e.g. :func:`conv`, :func:`pool`, :func:`reduce`, :func:`matmul`, ...).
+
+    See issue `#1691 <https://github.com/rwth-i6/returnn/issues/1691>`__.
+
+    :param default: what to return if it is not defined in the config,
+        and ``default_false_for_behavior_version_up_to`` does not apply.
+    :param default_false_for_behavior_version_up_to: if it is not defined in the config,
+        and if this is set, and the behavior version is less or equal,
+        then return False by default, i.e. do not use the mask by default, if it is not defined in the config.
+        This takes precedence over `default`.
+    :return: what to use for the ``use_mask`` argument by default
+    """
+    from returnn.config import get_global_config
+
+    config = get_global_config(raise_exception=False)
+    config_value = None
+    if config:
+        if "rf_use_mask" in config.typed_dict:
+            config_value = config.typed_dict["rf_use_mask"]
+            assert config_value is None or isinstance(config_value, bool)
+        elif "rf_use_mask" in config.dict:
+            config_value = config.bool("rf_use_mask", None)
+    if config_value is not None:
+        return config_value
+
+    if default_false_for_behavior_version_up_to is not None:
+        from returnn.util.basic import BehaviorVersion
+
+        if BehaviorVersion.get() <= default_false_for_behavior_version_up_to:
+            return False
+    return default

{returnn-1.20250226.115259 → returnn-1.20250226.183415}/returnn/frontend/normalization.py

@@ -218,10 +218,9 @@ class BatchNorm(rf.Module):
 
         if any(d.need_masking() for d in source.dims if d != self.in_dim):
             if self.use_mask is None:
-                raise ValueError(
-                    f"{self}: use_mask must be specified if the input {source} has any dynamic spatial dims"
-                )
-            use_mask = self.use_mask
+                use_mask = rf.use_mask_default(default=True)
+            else:
+                use_mask = self.use_mask
         else:
             use_mask = False  # not needed. False because this potentially enables an efficient fused op.
 

{returnn-1.20250226.115259 → returnn-1.20250226.183415}/returnn/frontend/signal.py

@@ -6,6 +6,7 @@ stft etc
 from __future__ import annotations
 from typing import Optional, Tuple
 from returnn.tensor import Tensor, Dim
+import returnn.frontend as rf
 
 
 __all__ = ["stft"]
@@ -23,6 +24,7 @@ def stft(
     window_enforce_even: bool = True,
     out_spatial_dim: Optional[Dim] = None,
     out_dim: Optional[Dim] = None,
+    use_mask: Optional[bool] = None,
 ) -> Tuple[Tensor, Dim, Dim]:
     """
     Calculate the short-time Fourier transform (STFT) of a signal.
@@ -65,8 +67,14 @@ def stft(
         but in most other frameworks, the behavior matches to window_enforce_even=False.
     :param out_spatial_dim:
     :param out_dim:
+    :param use_mask:
     :return: (stft, out_spatial_dim, out_dim)
     """
+    if in_spatial_dim.need_masking():
+        if use_mask is None:
+            use_mask = rf.use_mask_default(default=True, default_false_for_behavior_version_up_to=22)
+        if use_mask:
+            x = x.copy_masked(0, dims=[in_spatial_dim])
     fft_length = fft_length or frame_length
     if out_dim is None:
         out_dim = Dim(fft_length // 2 + 1, name="stft-freq")

{returnn-1.20250226.115259 → returnn-1.20250226.183415}/returnn/tensor/_dim_extra.py

@@ -422,14 +422,10 @@ class _DimMixin:
         :param func: operates inplace
         """
         dyn_size_ext = self.dyn_size_ext.copy() if self.dyn_size_ext is not None else None
-        dyn_size_ext_max = self._dyn_size_max_value if self._dyn_size_max_value is not None else None
         self.reset_raw(only_self=True)
         if dyn_size_ext is not None:
             func(dyn_size_ext)
-        if dyn_size_ext_max is not None:
-            func(dyn_size_ext_max)
         self.dyn_size_ext = dyn_size_ext
-        self._dyn_size_max_value = dyn_size_ext_max
 
     def _can_use_in_ctx(self, ctx):
         """

{returnn-1.20250226.115259 → returnn-1.20250226.183415}/returnn/tensor/_tensor_extra.py

@@ -2966,7 +2966,7 @@ class _TensorMixin(_TensorMixinBase):
         mask = None
         for axis in axes:
             mask_ = self._dims[axis].get_mask(dim_order=self.dims, device=self.device)
-            mask = rf.logical_and(mask, mask_) if mask is not None else mask_
+            mask = rf.combine_bc(mask, "logical_and", mask_) if mask is not None else mask_
         assert isinstance(mask, _t.Tensor)
         res = rf.where(mask, self, mask_value)
         if use_padding_info:

{returnn-1.20250226.115259 → returnn-1.20250226.183415}/returnn/tensor/tensor_dict.py

@@ -9,7 +9,7 @@ We also might have model_outputs in the user config.
 """
 
 from __future__ import annotations
-from typing import Optional, Union, Any, Type, Dict, Sequence
+from typing import Optional, Union, Any, Type, Dict, Sequence, List
 from .tensor import Tensor
 from .dim import Dim
 
@@ -160,6 +160,20 @@ class TensorDict:
                         assert dim.size == raw_tensor_dict[key_]
                 visited_dims.add(dim)
 
+    def all_dims(self) -> List[Dim]:
+        """
+        :return: list of dims
+        """
+        visited_dims = set()
+        out = []
+        for key, value in self.data.items():
+            for dim in value.dims:
+                if dim in visited_dims:
+                    continue
+                out.append(dim)
+                visited_dims.add(dim)
+        return out
+
 
 def _convert_to_tensor(opts: _TensorT, *, name: Optional[str] = None) -> Tensor:
     """

{returnn-1.20250226.115259 → returnn-1.20250226.183415}/returnn/util/basic.py

@@ -219,7 +219,7 @@ class BehaviorVersion:
     See :ref:`behavior_version`.
     """
 
-    _latest_behavior_version = 22
+    _latest_behavior_version = 23
     _behavior_version = None  # type: typing.Optional[int]
     _min_behavior_version = 0  # type: int
 

{returnn-1.20250226.115259 → returnn-1.20250226.183415/returnn.egg-info}/PKG-INFO

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

{returnn-1.20250226.115259 → returnn-1.20250226.183415}/tests/rf_utils.py

@@ -13,7 +13,7 @@ import numpy.testing
 from returnn.config import Config, global_config_ctx
 from returnn.util.pprint import pprint
 import returnn.frontend as rf
-from returnn.tensor import Tensor, Dim, TensorDict
+from returnn.tensor import Tensor, Dim, TensorDict, batch_dim
 from returnn.tensor.utils import tensor_dict_fill_random_numpy_
 import returnn.torch.frontend as rft
 from returnn.torch.data.tensor_utils import tensor_dict_numpy_to_torch_, tensor_dict_torch_to_numpy_
@@ -57,9 +57,11 @@ def run_model(
     dyn_dim_min_sizes: Optional[Dict[Dim, int]] = None,
     test_tensorflow: bool = True,
     allow_inf_nan_in_output: bool = False,
+    test_single_batch_entry: bool = False,  # can later enable this globally
 ) -> TensorDict:
     """run"""
     print(f"* run_model with dyn_dim_max_sizes={dyn_dim_max_sizes!r}")
+    extern_data_dims = extern_data.all_dims()
     extern_data.reset_content()
     tensor_dict_fill_random_numpy_(
         extern_data, dyn_dim_max_sizes=dyn_dim_max_sizes, dyn_dim_min_sizes=dyn_dim_min_sizes
@@ -85,6 +87,21 @@ def run_model(
             )
             raise Exception(f"Non-finite values in output: {non_finite_outputs}. See log above.")
 
+    if test_single_batch_entry and batch_dim in extern_data_dims:
+        dyn_dims = [
+            d
+            for d in extern_data_dims
+            if d.dyn_size_ext is not None
+            and d.dyn_size_ext.dims == (batch_dim,)
+            and d.dyn_size_ext.raw_tensor.min() != d.dyn_size_ext.raw_tensor.max()
+        ]
+        if dyn_dims:  # e.g. the typical time dim with dyn size shape [batch_dim]
+            batch_idx = dyn_dims[0].dyn_size_ext.raw_tensor.argmin().item()
+            print(f"** run with PyTorch backend with single batch entry for some smaller sequence {batch_idx=}")
+            for d in dyn_dims:
+                print(f"  {d}: {d.dyn_size_ext.raw_tensor}")
+            _run_model_torch_single_batch(extern_data, get_model, forward_step, batch_idx=batch_idx, ref_output=out_pt)
+
     if not test_tensorflow:
         return out_pt
     if disable_tf:
@@ -125,6 +142,7 @@ def run_model(
     assert set(out_pt_raw.keys()) == set(out_tf_raw.keys())
     for k, v_pt in out_pt_raw.items():
         v_tf = out_tf_raw[k]
+        print(f"  comparing {k!r} {_array_repr(v_pt)} PT vs TF")
         numpy.testing.assert_allclose(v_pt, v_tf, atol=1e-5, rtol=1e-5, err_msg=f"output {k!r} differs")
     return out_pt
 
@@ -222,6 +240,72 @@ def run_model_torch_train(
     return res
 
 
+def _run_model_torch_single_batch(
+    extern_data: TensorDict,
+    get_model: rf.GetModelFunc,
+    forward_step: rf.StepFunc,
+    *,
+    batch_idx: int,
+    ref_output: TensorDict,
+):
+    """
+    Restrict batch_dim to batch_idx only.
+    This is somewhat hacky:
+    We want to keep the same dim tags.
+    Thus, we overwrite all the data by the sliced data.
+    We also want to keep the batch_dim for the forward_step as the user code probably expects it.
+    Afterward, we want to recover the original data.
+    """
+    # Store original data to be able to recover it later.
+    extern_data_raw = extern_data.as_raw_tensor_dict(expected_value_type=numpy.ndarray)
+
+    # noinspection PyShadowingNames
+    def _get_slices(x: Tensor) -> Tuple[slice, ...]:
+        slices = []
+        for dim in x.dims:
+            if dim == batch_dim:
+                slices.append(slice(batch_idx, batch_idx + 1))
+            elif dim.dyn_size_ext is not None and batch_dim in dim.dyn_size_ext.dims:
+                slices.append(slice(0, dim.dyn_size_ext.raw_tensor.max().item()))
+            else:
+                slices.append(slice(None))
+        return tuple(slices)
+
+    # inplace
+    # noinspection PyShadowingNames
+    def tensor_numpy_restrict_batch_dim_(x: Tensor[numpy.ndarray]):
+        if batch_dim not in x.dims:
+            return
+        if x.raw_tensor is not None:
+            assert isinstance(x.raw_tensor, numpy.ndarray)
+            x.raw_tensor = x.raw_tensor[_get_slices(x)]
+
+    batch_dim.reset_eager()
+    batch_dim.dyn_size_ext = Tensor("batch_size", dims=(), dtype="int32", raw_tensor=numpy.array(1, dtype="int32"))
+    for dim in extern_data.all_dims():
+        dim.transform_tensors(tensor_numpy_restrict_batch_dim_)
+    for v in extern_data.data.values():
+        tensor_numpy_restrict_batch_dim_(v)
+
+    output = _run_model_torch(extern_data, get_model, forward_step)
+    for key, ref_output_ in ref_output.data.items():
+        output_ = output.data[key]
+        if batch_dim not in ref_output_.dims:
+            continue
+        batch_axis = ref_output_.dims.index(batch_dim)
+        assert output_.dims[batch_axis] == batch_dim and len(output_.dims) == len(ref_output_.dims)
+        # Slice the raw ref output to be able to match it to the raw single output.
+        ref_output_raw = ref_output_.raw_tensor[_get_slices(output_)]
+        single_output_raw = output_.raw_tensor
+        numpy.testing.assert_allclose(
+            ref_output_raw, single_output_raw, atol=1e-5, rtol=1e-5, err_msg=f"output {key!r} differs"
+        )
+
+    # Recover original data.
+    extern_data.reset_content()
+    extern_data.assign_from_raw_tensor_dict_(extern_data_raw)
+
+
 def _run_model_net_dict_tf(
     extern_data: TensorDict, get_model: rf.GetModelFunc, forward_step: rf.StepFunc
 ) -> TensorDict:
@@ -349,15 +433,9 @@ def _pad_mask_zeros(x: Union[TensorDict, Tensor, Dim]):
         return
 
     assert isinstance(x, Tensor)
-    for i, d in enumerate(x.dims):
+    x.raw_tensor = x.copy_masked(0).raw_tensor
+    for d in x.dims:
         _pad_mask_zeros(d)
-        if d.need_masking():
-            mask = x.get_sequence_mask_tensor(i)
-            if not set(mask.dims).issubset(set(x.dims)):
-                print(f"Warning: cannot apply mask {mask} for dim {d} on tensor {x}.")
-                continue
-            mask_raw = mask.copy_compatible_to_dims_raw(x.dims)
-            x.raw_tensor = numpy.where(mask_raw, x.raw_tensor, numpy.zeros((), dtype=x.raw_tensor.dtype))
 
 
 def _check_dim(d_pt: Dim, d_tf: Dim):
@@ -403,3 +481,17 @@ def _walk_dims(start: Dim, *, func=print):
             for k, v in dim_extra.same_for_batch_ctx.items():
                 k: Any
                 queue.append((path + ("_extra.same_for_batch_ctx", k), v))
+
+
+def _array_repr(x: Union[numpy.ndarray, numpy.number]) -> str:
+    if not isinstance(x, numpy.ndarray):
+        return f"<{type(x).__name__} {x!r}>"
+
+    try:
+        import lovely_numpy
+
+        return f"<{lovely_numpy.lovely(x)}>"
+    except ImportError:
+        if x.size <= 10:
+            return repr(x)
+        return f"<array shape={x.shape} dtype={x.dtype} min={x.min()} max={x.max()}>"

{returnn-1.20250226.115259 → returnn-1.20250226.183415}/tests/test_rf_array.py

@@ -193,9 +193,16 @@ def test_pad_packed_batched():
     in_ = out_dict["in"]
     flat = out_dict["flat"]
     out = out_dict["out"]
-    print(in_.raw_tensor.shape, flat.raw_tensor.shape, out.raw_tensor.shape)
+    print("in:", in_, in_.raw_tensor.shape)
+    print("in time1:", in_.dims[1].dyn_size)
+    print("in time2:", in_.dims[2].dyn_size)
+    print("flat:", flat, flat.raw_tensor.shape)
+    print("out:", out, out.raw_tensor.shape)
+    print("in raw:")
     print(in_.raw_tensor)
+    print("flat raw:")
     print(flat.raw_tensor)
+    print("out raw:")
     print(out.raw_tensor)
     np.testing.assert_array_equal(in_.raw_tensor, out.raw_tensor)
 

{returnn-1.20250226.115259 → returnn-1.20250226.183415}/tests/test_rf_encoder_conformer.py

@@ -78,6 +78,7 @@ def test_e_branchformer():
     import torch
     import returnn.frontend as rf
     from returnn.util.debug import PyTracer, check_py_traces_rf_to_pt_equal
+    from returnn.config import global_config_ctx, Config
 
     rf.select_backend_torch()
     rf.set_random_seed(42)
@@ -300,7 +301,9 @@ def test_e_branchformer():
         ],
         Tensor,
     ) as trace_rf, torch.no_grad():
-        enc_out, _ = model_rf(enc_in, in_spatial_dim=enc_spatial_dim)
+        # ESPnet E-Branchformer does not use masking properly. Keep it disabled here as well.
+        with global_config_ctx(Config({"rf_use_mask": False})):
+            enc_out, _ = model_rf(enc_in, in_spatial_dim=enc_spatial_dim)
     enc_out = enc_out.copy_transpose((batch_dim, enc_spatial_dim, model_dim))
     enc_out = enc_out.copy_masked(0.0)
 

returnn-1.20250226.115259/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20250226.115259'
-long_version = '1.20250226.115259+git.0d32534'