returnn 1.20250226.104737__tar.gz → 1.20250226.132109__tar.gz

{returnn-1.20250226.104737/returnn.egg-info → returnn-1.20250226.132109}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20250226.104737
+Version: 1.20250226.132109
 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.132109/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20250226.132109'
+long_version = '1.20250226.132109+git.7dbfd3a'

{returnn-1.20250226.104737 → returnn-1.20250226.132109}/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.104737 → returnn-1.20250226.132109}/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.104737 → returnn-1.20250226.132109}/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.104737 → returnn-1.20250226.132109}/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.104737 → returnn-1.20250226.132109}/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.104737 → returnn-1.20250226.132109}/returnn/torch/util/debug_inf_nan.py

@@ -39,6 +39,7 @@ from io import TextIOBase
 import traceback
 from types import FrameType
 import torch
+import tree
 
 # noinspection PyProtectedMember
 from torch.utils._python_dispatch import TorchDispatchMode
@@ -96,6 +97,7 @@ def debug_inf_nan(
 
 # For efficiency, and to be less spammy
 _TraceFuncNameBlacklist = {
+    "aten::empty.memory_format",
     "aten::zeros_like",
     "aten::ones_like",
     "aten::full",
@@ -113,6 +115,7 @@ _TraceFuncNameBlacklist = {
     "aten::split_with_sizes",
     "aten::slice.Tensor",
     "aten::select.int",
+    "aten::max_pool2d_with_indices",
 }
 
 
@@ -140,19 +143,20 @@ class _TraceOps(TorchDispatchMode):
         if self.report_every_op_call:
             print(f"--- op {func.name()}", file=self.file)
         out = func(*args, **kwargs)
-        if isinstance(out, torch.Tensor):
-            with no_python_dispatcher():
-                got_nan_inf_t = torch.stack([torch.isnan(out).any(), torch.isinf(out).any()]).cpu()
-                got_nan = got_nan_inf_t[0].item()
-                got_inf = got_nan_inf_t[1].item()
-                if got_nan or got_inf:
-                    s = "/".join([s_ for s_, b in [("nan", got_nan), ("inf", got_inf)] if b])
-                    print(f"--> {s} in {func}: {out}", file=self.file)
-                    traceback.print_list(
-                        _extract_stack_up_to(skip_top_num_frames=1, root_frame=self.root_frame), file=self.file
-                    )
-                    if self.stop_reporting_after_first_inf_nan:
-                        self.enabled = False
+        for out_ in tree.flatten(out):
+            if isinstance(out_, torch.Tensor):
+                with no_python_dispatcher():
+                    got_nan_inf_t = torch.stack([torch.isnan(out_).any(), torch.isinf(out_).any()]).cpu()
+                    got_nan = got_nan_inf_t[0].item()
+                    got_inf = got_nan_inf_t[1].item()
+                    if got_nan or got_inf:
+                        s = "/".join([s_ for s_, b in [("nan", got_nan), ("inf", got_inf)] if b])
+                        print(f"--> {s} in {func}: {out_}", file=self.file)
+                        traceback.print_list(
+                            _extract_stack_up_to(skip_top_num_frames=1, root_frame=self.root_frame), file=self.file
+                        )
+                        if self.stop_reporting_after_first_inf_nan:
+                            self.enabled = False
         return out
 
 

{returnn-1.20250226.104737 → returnn-1.20250226.132109}/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.104737 → returnn-1.20250226.132109/returnn.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20250226.104737
+Version: 1.20250226.132109
 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.104737 → returnn-1.20250226.132109}/tests/rf_utils.py

@@ -17,6 +17,7 @@ from returnn.tensor import Tensor, Dim, TensorDict
 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_
+from returnn.torch.util.debug_inf_nan import debug_inf_nan as torch_debug_inf_nan
 
 # noinspection PyProtectedMember
 from returnn.frontend._random_journal import RandomJournal
@@ -72,8 +73,17 @@ def run_model(
         out_pt_raw = out_pt.as_raw_tensor_dict(include_const_sizes=True)
 
     if not allow_inf_nan_in_output:
+        non_finite_outputs = {}
         for k, v in out_pt.data.items():
-            assert numpy.isfinite(v.raw_tensor).all(), f"output {k!r} has non-finite values: {v.raw_tensor}"
+            if not numpy.isfinite(v.raw_tensor).all():
+                non_finite_outputs[k] = v
+                print(f"ERROR: output {k!r} has non-finite values:\n{v.raw_tensor}")
+        if non_finite_outputs:
+            torch_debug_inf_nan(
+                lambda: (_run_model_torch(extern_data, get_model, forward_step), None)[-1],
+                stop_reporting_after_first_inf_nan=False,
+            )
+            raise Exception(f"Non-finite values in output: {non_finite_outputs}. See log above.")
 
     if not test_tensorflow:
         return out_pt

{returnn-1.20250226.104737 → returnn-1.20250226.132109}/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.104737 → returnn-1.20250226.132109}/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.104737/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20250226.104737'
-long_version = '1.20250226.104737+git.5a134ff'