returnn 1.20250226.115259__py3-none-any.whl → 1.20250226.132109__py3-none-any.whl

returnn/PKG-INFO

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

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

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/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/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/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/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/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.dist-info → returnn-1.20250226.132109.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20250226.115259
+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.115259.dist-info → returnn-1.20250226.132109.dist-info}/RECORD

@@ -1,9 +1,9 @@
-returnn/PKG-INFO,sha256=2Ws--V5aicc3WJ-I6OrqPbbrvVNTH3Cnno6L7yeIyKY,5215
+returnn/PKG-INFO,sha256=1tdi71kdx2tM5iCiEhtBkyM0jO74aiySG7bvbniZgz0,5215
 returnn/__init__.py,sha256=biBtRsM0WZ406vShaeH-9WFoqJ8XwTbn6g0EeFJ7l8E,1012
 returnn/__main__.py,sha256=qBFbuB1yN3adgVM5pXt2-Yq9vorjRNchNPL8kDKx44M,31752
 returnn/__old_mod_loader__.py,sha256=nvsNY-xELdS_IPNkv66Q9Rmvg4dbGW0-EBRDcCmctos,7654
 returnn/__setup__.py,sha256=22kQn2fh11iPM0hLb2Fy5sLmoU1JGvmDxXRYuRgQkwU,4659
-returnn/_setup_info_generated.py,sha256=3ur2a8rg2h6MJe2vAo7Tq4axfkV1GYJMcaQdnsmshb8,77
+returnn/_setup_info_generated.py,sha256=GFgn-7LeKnvijLbLCBASDWIfBb96hZ17jlTud_Q9xAM,77
 returnn/config.py,sha256=3tmKhB6FnQZaNdtcYsiB61JnEY--iZ2qmJ4yq0b6tE0,29140
 returnn/forward_iface.py,sha256=A_OJiaXsX4MlXQRzST86ylyxSUZbC402PQL1REcqHjM,911
 returnn/learning_rate_control.py,sha256=ZvWryAn_tv9DhV8sh1LV3eE34Yltl3On3mYZAG4hR9s,34684
@@ -80,7 +80,7 @@ returnn/frontend/_cache.py,sha256=JAhi7L-raQ3A-NC3JUYDtdRTwT3BGJJGGZxrZ8MfEWQ,84
 returnn/frontend/_numpy_backend.py,sha256=2oCtG0YCWL_89v4cD_jDj8em1O_Fp-_YWl5EblGi_yo,7858
 returnn/frontend/_random_journal.py,sha256=_ktP_mjgx8vtQQGX_DofdhewJj0aPiczefTWeemPkmo,5457
 returnn/frontend/_utils.py,sha256=4A3MSRM0i86J77550uR_AjcBEPu6nymLUZ9Xd1V3Fkc,12073
-returnn/frontend/array_.py,sha256=M5vCeH0nlwJ-zrdjbZpsMLN6StOLn0iM7PnXvGLLE3g,49154
+returnn/frontend/array_.py,sha256=eYwH-NVAoHpVrFdJv08lCqh3jvfoZV_ZBEoWHjsBz0o,50090
 returnn/frontend/attention.py,sha256=GKt-Xqnz8sIyXVrE0i4VCS7J2Wu7dmoH_BA0Cu8CrXQ,45769
 returnn/frontend/backend.py,sha256=iQ9w4xl8Ea7bgpb0VUaCKq50rV5Bl2E5J8Rhd-oqD_c,883
 returnn/frontend/build_from_dict.py,sha256=rfWa2rjjhIR_kIQED_nMrygrQBunS6unegzWTLVbC98,3017
@@ -88,9 +88,9 @@ returnn/frontend/cond.py,sha256=gh6wg0aSbAJQfKRv4BQAu-EfPWtWPLFjgc8IaPPFmwg,1023
 returnn/frontend/const.py,sha256=bL51HXxq858dWmrKd61k8tWBWIe67jVf9pj1wZcZZAo,3945
 returnn/frontend/container.py,sha256=wF3OlQN7WlOVmmdapUth_Unha3DVf6h1B7okBJAuJDA,8011
 returnn/frontend/control_flow_ctx.py,sha256=v17CsNwRnZYe8GdMtGJt2ftibfxMCGK1i0l-GX5ILu0,699
-returnn/frontend/conv.py,sha256=51LZovcRzITDLXvPcJs_MFsGEY_MFvO_MFF9D-jZstA,22481
+returnn/frontend/conv.py,sha256=p4R6j40GCvVrw3kbQQJtfxY6tfIR8Rb3tIzwAtiLuec,23858
 returnn/frontend/device.py,sha256=K7Y1qoQcO4GIHgLkPLQWK-GVT8gKL8GwyQrmPo8LgBE,1438
-returnn/frontend/dims.py,sha256=hKA7IQRB0DbohN1ngNw31W44BsyjdHCtYAccxOcumzQ,10872
+returnn/frontend/dims.py,sha256=aH5FQ_m0xMD6Rj-BUWGx8lB-HkCuwZfMBf6mZbGGW5E,12611
 returnn/frontend/dropout.py,sha256=rsx3p5b0NblBfXXSQZTQFJ8jUUS3fj4Qzc39iffBMCA,5006
 returnn/frontend/dtype.py,sha256=Ooc5BrcNrTp6XShuFEV9g5V6-niuy4ImP_Lt_Qgq3jE,1886
 returnn/frontend/gradient.py,sha256=G-Qv4gKGHYEeB92Zwco9ao4qjd6umZPUzQC4J-fbYWo,4033
@@ -105,7 +105,7 @@ returnn/frontend/math_.py,sha256=KlJxdIib8ENlid7cc4lcwHv5e21tzTjTEV8VgEDAijo,169
 returnn/frontend/matmul.py,sha256=3QaGiZtSs9PriT40T7Vc3KnYKPgYSN4tCZytYeq9qMA,1945
 returnn/frontend/module.py,sha256=219rh5mE0CD0-NdxXLsKyhv3BNtOI9jSyiI1Rb8MOyU,10700
 returnn/frontend/nested.py,sha256=CT3C0wXkeWGjJcAoF6yebsRXuN8-YpjO2eqgdl1-vaE,11005
-returnn/frontend/normalization.py,sha256=fvkrMB0xlBeeagFHDGnZ7SvgpJ2TBirfhq3u0RGJFHw,14139
+returnn/frontend/normalization.py,sha256=QIjXYg0C8BD2g_1lAkVO4Cara729uHC_bsQh99VsWeI,14061
 returnn/frontend/parameter.py,sha256=w6SN-uv87OyeWBt90_3UBbK0h6sftSOCxkqXPg76caY,10375
 returnn/frontend/parametrizations.py,sha256=hVbOlgm1pQAmZnAnNxq8Tk23rykr_iy3-6R1H6CwlMA,2798
 returnn/frontend/parametrize.py,sha256=VhgTEP7ehON950Q4bkCy8rvg9641moEKAXn0XzomK6E,7216
@@ -114,7 +114,7 @@ returnn/frontend/rand.py,sha256=Levgf5VtOOBKDSgz0869Jf3VW4BWxYZuRXsa_fOxNI4,1296
 returnn/frontend/rec.py,sha256=4m20LvsPJ75pRYykVrup6Csj_D7duG-dW28SaJh-sq8,7863
 returnn/frontend/reduce.py,sha256=-Zt-OH6Zbtb9uR6YEzurCyrowH-anIXvuga6Pla2V70,10220
 returnn/frontend/run_ctx.py,sha256=ItcZwuFItkZjYWrg715L1Za2Xg7__MQCrRCAwBeTUxA,21411
-returnn/frontend/signal.py,sha256=8miFjBropGt2xQNwjLcSd9VaiKc04nHo04iPfEwdB8M,4078
+returnn/frontend/signal.py,sha256=XgOBL1iy-cJgulePH5HRPAwp2cScy60q4RItr7xzvGc,4412
 returnn/frontend/state.py,sha256=EePdrx6PtWL4mJ2XZmGlh5dl4nq6G9wZpqP4hdDEzfY,2935
 returnn/frontend/stepwise_scheduler.py,sha256=fMOTR7npGCDXrXDmSQ4VwmudoHEbY3Yr-QGyjFdQJSc,927
 returnn/frontend/tensor_array.py,sha256=Ej7CHtvpY0yBROlAk5vFe3CTXh-iAuqu9qcXS3Qxt2I,4328
@@ -233,7 +233,7 @@ returnn/torch/util/gradient_checkpoint.py,sha256=iLy-FB65DC8O6LxzmMvFjnSdpIVpko8
 returnn/torch/util/module.py,sha256=MXHIrF9Isu575DDJIa81212ULKwdqu1oOLxDVZecVSk,1693
 returnn/torch/util/scaled_gradient.py,sha256=3585VuNypBty-pW6r3BKK047H3MqZQSdMjXeYAb4cmU,3192
 returnn/util/__init__.py,sha256=UIG1qw4idqhW71BV60ha7h9PktxvEVcBIu0lYRossK8,336
-returnn/util/basic.py,sha256=__rtDp8crZfm0mEeAKsRxNCdWuBHh9OeOm8UO-X4CJU,142380
+returnn/util/basic.py,sha256=Iynt9ATEs_8DaZsX5z6weMyaO2xW9o3gaywq6X7mbEc,142380
 returnn/util/better_exchook.py,sha256=MVMnuu6KoyqgvlMeQLQNTfdspcPR9MwigCXOpeTVqCI,62956
 returnn/util/bpe.py,sha256=LWFhICZsEOnMwNws0lybPNzKRX6rSr8yKCvP65vjl9Y,19656
 returnn/util/debug.py,sha256=wuRzdg9zB84WWCGyTjmRR_zYypu8gXxlc0nZ6si9OC8,28224
@@ -253,8 +253,8 @@ returnn/util/sig_proc.py,sha256=Tjz0VOAVyqu2qDCF5HZ1JjALjcFsHcNkcd96WgZeKfE,7265
 returnn/util/task_system.py,sha256=y4sMVXQ25Qd2z0rx03uOlXlkE-jbCYC1Sjfn-XlraVU,26003
 returnn/util/train_proc_manager.py,sha256=Pjht28k6uz6BNQ47uW6Gf880iyq5q4wx7P_K2tmoAM8,3266
 returnn/util/watch_memory.py,sha256=BR5P2kvBN6UI81cE0_1WAA6Hd1SByLbBaiDxvLhPOew,4213
-returnn-1.20250226.115259.dist-info/LICENSE,sha256=ywBD_U2aD4vpuoIgNAsjIGBYydl0tVKll3De0Z8s77c,11041
-returnn-1.20250226.115259.dist-info/METADATA,sha256=2Ws--V5aicc3WJ-I6OrqPbbrvVNTH3Cnno6L7yeIyKY,5215
-returnn-1.20250226.115259.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-returnn-1.20250226.115259.dist-info/top_level.txt,sha256=Lsn4WZc5Pbfk0-xDQOgnFCxOoqxL4CyeM3N1TFbJncw,8
-returnn-1.20250226.115259.dist-info/RECORD,,
+returnn-1.20250226.132109.dist-info/LICENSE,sha256=ywBD_U2aD4vpuoIgNAsjIGBYydl0tVKll3De0Z8s77c,11041
+returnn-1.20250226.132109.dist-info/METADATA,sha256=1tdi71kdx2tM5iCiEhtBkyM0jO74aiySG7bvbniZgz0,5215
+returnn-1.20250226.132109.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+returnn-1.20250226.132109.dist-info/top_level.txt,sha256=Lsn4WZc5Pbfk0-xDQOgnFCxOoqxL4CyeM3N1TFbJncw,8
+returnn-1.20250226.132109.dist-info/RECORD,,