returnn 1.20250304.10039__py3-none-any.whl → 1.20250304.113330__py3-none-any.whl

returnn/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20250304.10039
+Version: 1.20250304.113330
 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.20250304.010039'
-long_version = '1.20250304.010039+git.3e53d74'
+version = '1.20250304.113330'
+long_version = '1.20250304.113330+git.acf09da'

returnn/frontend/_backend.py

@@ -1223,7 +1223,7 @@ class Backend(Generic[T]):
         out_spatial_dims: Optional[Sequence[Dim]] = None,
         filter: Tensor,
         filter_size: Sequence[Dim],  # to have the order well-defined
-        padding: str,
+        padding: Union[str, int, Sequence[int]],
         strides: Optional[Union[int, Sequence[int]]] = None,
         dilation_rate: Optional[Union[int, Sequence[int]]] = None,
         groups: Optional[int] = None,
@@ -1258,7 +1258,7 @@ class Backend(Generic[T]):
         *,
         mode: str,
         pool_size: Sequence[int],
-        padding: str = "valid",
+        padding: Union[str, int, Sequence[int]] = "valid",
         dilation_rate: Union[Sequence[int], int] = 1,
         strides: Sequence[int],
         in_spatial_dims: Sequence[Dim],

returnn/frontend/conv.py

@@ -181,15 +181,46 @@ def conv(
     in_spatial_dims: Sequence[Dim],
     out_spatial_dims: Optional[Sequence[Dim]] = None,
     filter: Tensor,
-    filter_size: Sequence[Dim],  # to have the order well-defined
-    padding: str,
+    filter_size: Sequence[Dim],
+    padding: Union[str, int, Sequence[int]],
     strides: Optional[Union[int, Sequence[int]]] = None,
     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"""
+    """
+    Generic N-D convolution.
+
+    :param source:
+    :param in_dim: input channels
+    :param out_dim: output channels
+    :param in_spatial_dims: On what dimensions to operate on.
+        The number of specified dims (1, 2 or 3) specifies whether this is 1D, 2D or 3D convolution.
+        The order is consistent with the order of the ``filter_size``, ``strides``, etc.
+    :param out_spatial_dims:
+    :param filter:
+    :param filter_size: defines the order of dims in ``filter``
+        such that it matches the order of ``in_spatial_dims``.
+    :param padding: "valid" or "same" or int. "valid" is like padding=0.
+        padding="same" will pad such that the output has the same spatial dimensions as the input
+        (in case of stride=1), or otherwise ceildiv(input, stride).
+        The specific padding in padding="same" with stride>1 has changed with behavior version >=24
+        (or global config option ``rf_use_consistent_same_padding``)
+        and is now consistent independent of dimension size.
+        See :func:`_consistent_same_padding` for more details.
+    :param strides: the default (if it is None) is 1
+    :param dilation_rate:
+    :param groups:
+    :param bias:
+    :param use_mask: Whether to mask the input tensor based on seq lengths
+        such that the padding in the padded tensor is ignored
+        (it will mask with 0).
+        With behavior version >=23, this is enabled by default,
+        or configured with global config option ``rf_use_mask``.
+        (Also see :func:`use_mask_default`).
+    :return: out, out_spatial_dims
+    """
     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)
@@ -198,6 +229,10 @@ def conv(
     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}")
+    if padding == "same" and _any_is_non_default(strides, default=1) and _should_use_consistent_same_padding():
+        source, in_spatial_dims, padding = _consistent_same_padding(
+            source, in_spatial_dims=in_spatial_dims, filter_size=filter_size, dilation_rate=dilation_rate, pad_value=0
+        )
     # noinspection PyProtectedMember
     out, out_spatial_dims = source._raw_backend.conv(
         source,
@@ -359,6 +394,9 @@ def transposed_conv(
             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)
+    if padding == "same" and _any_is_non_default(strides, default=1) and _should_use_consistent_same_padding():
+        # I don't really know what this should mean here... Investigate this further...
+        raise NotImplementedError("consistent same padding not implemented for transposed conv")
     # noinspection PyProtectedMember
     out, out_spatial_dims = source._raw_backend.transposed_conv(
         source=source,
@@ -409,7 +447,7 @@ def pool(
     nd: Optional[int] = None,
     mode: str,
     pool_size: Union[Sequence[int], int],
-    padding: str = "valid",
+    padding: Union[str, int, Sequence[int]] = "valid",
     dilation_rate: Union[Sequence[int], int] = 1,
     strides: Optional[Union[Sequence[int], int]] = None,
     in_spatial_dims: Union[Sequence[Dim], Dim],
@@ -417,19 +455,29 @@ def pool(
     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.
+    Generic N-D pooling.
 
     :param source:
     :param nd:
     :param mode: "max" or "avg"
     :param pool_size: shape of the window of each reduce
-    :param padding: "valid" or "same"
+    :param padding: "valid" or "same" or int. "valid" is like padding=0.
+        padding="same" will pad such that the output has the same spatial dimensions as the input
+        (in case of stride=1), or otherwise ceildiv(input, stride).
+        The specific padding in padding="same" with stride>1 has changed with behavior version >=24
+        (or global config option ``rf_use_consistent_same_padding``)
+        and is now consistent independent of dimension size.
+        See :func:`_consistent_same_padding` for more details.
     :param dilation_rate:
-    :param strides: in contrast to tf.nn.pool, the default (if it is None) will be set to pool_size
+    :param strides: the default (if it is None) will be set to pool_size (in contrast to :func:`conv`)
     :param in_spatial_dims:
     :param out_spatial_dims:
-    :param use_mask:
+    :param use_mask: Whether to mask the input tensor based on seq lengths
+        such that the padding in the padded tensor is ignored
+        (for max-pooling, it will mask with -inf, for avg-pooling with 0).
+        With behavior version >=23, this is enabled by default,
+        or configured with global config option ``rf_use_mask``.
+        (Also see :func:`use_mask_default`).
     :return: out, out_spatial_dims
     """
     if isinstance(in_spatial_dims, Dim):
@@ -451,8 +499,7 @@ def pool(
         strides = pool_size
     elif isinstance(strides, int):
         strides = [strides] * nd
-    assert isinstance(strides, (list, tuple))
-    assert len(strides) == nd
+    assert isinstance(strides, (list, tuple)) and len(strides) == nd and all(isinstance(s, int) for s in strides)
 
     if any(in_spatial_dim.need_masking() for in_spatial_dim in in_spatial_dims):
         if use_mask is None:
@@ -462,6 +509,15 @@ def pool(
     else:
         use_mask = False
 
+    if padding == "same" and _any_is_non_default(strides, default=1) and _should_use_consistent_same_padding():
+        source, in_spatial_dims, padding = _consistent_same_padding(
+            source,
+            in_spatial_dims=in_spatial_dims,
+            filter_size=pool_size,
+            dilation_rate=dilation_rate,
+            pad_value={"max": float("-inf"), "avg": 0}[mode],
+        )
+
     # noinspection PyProtectedMember
     out, out_spatial_dims = source._raw_backend.pool(
         source=source,
@@ -642,7 +698,7 @@ def make_conv_out_spatial_dims(
     in_spatial_dims: Sequence[Dim],
     *,
     filter_size: Union[Sequence[Union[int, Dim]], int, Dim],
-    padding: str,
+    padding: Union[str, int, Sequence[int]],
     strides: Union[Sequence[int], int] = 1,
     dilation_rate: Union[Sequence[int], int] = 1,
     description_prefix: Optional[str] = None,
@@ -658,11 +714,15 @@ def make_conv_out_spatial_dims(
     if isinstance(dilation_rate, int):
         dilation_rate = [dilation_rate] * nd
     assert nd == len(in_spatial_dims) == len(filter_size) == len(strides) == len(dilation_rate)
-    assert padding.lower() in ("valid", "same")
+    if isinstance(padding, (int, str)):
+        padding = [padding] * nd
+    padding = [p.lower() if isinstance(p, str) else p for p in padding]
     out_spatial_dims = []
     for i in range(nd):
         in_spatial_dim = in_spatial_dims[i]
-        if filter_size[i] == strides[i] == 1 or (strides[i] == 1 and padding.lower() == "same"):
+        if (filter_size[i] == strides[i] == 1 and padding[i] in ("valid", "same", 0)) or (
+            strides[i] == 1 and padding[i] == "same"
+        ):
             out_spatial_dims.append(in_spatial_dim)
         else:
             out_spatial_dim = _calc_out_dim(
@@ -670,7 +730,7 @@ def make_conv_out_spatial_dims(
                 filter_size=filter_size[i],
                 stride=strides[i],
                 dilation_rate=dilation_rate[i],
-                padding=padding,
+                padding=padding[i],
             )
             assert isinstance(out_spatial_dim, Dim)
             if description_prefix and out_spatial_dim != in_spatial_dim:
@@ -681,7 +741,7 @@ def make_conv_out_spatial_dims(
                     filter_size=filter_size[i],
                     stride=strides[i],
                     dilation_rate=dilation_rate[i],
-                    padding=padding,
+                    padding=padding[i],
                 )
             out_spatial_dims.append(out_spatial_dim)
     return out_spatial_dims
@@ -695,7 +755,7 @@ def _calc_out_dim(in_dim, filter_size, stride, padding, dilation_rate=1):
     :param int filter_size: e.g. 2, for the corresponding axis
     :param int stride: e.g. 1, for the corresponding axis
     :param int dilation_rate: e.g. 1
-    :param str padding: "valid" or "same"
+    :param str|int padding: "valid" or "same" or int
     :return: the output dimension
     :rtype: T
     """
@@ -712,13 +772,16 @@ def _calc_out_dim(in_dim, filter_size, stride, padding, dilation_rate=1):
             return rf.ceil_divide(a, b)
         return -(-a // b)
 
-    padding = padding.upper()
+    padding = padding.lower() if isinstance(padding, str) else padding
     # See tf.compat.v1.nn.convolution() documentation for more.
-    if padding == "SAME":
+    if padding == "same":
         if isinstance(in_dim, Dim):
             return in_dim.ceildiv_right(stride)
         return ceildiv(in_dim, stride)
-    elif padding == "VALID":
+    elif padding == "valid" or isinstance(padding, int):
+        if isinstance(padding, int) and padding != 0:
+            assert padding > 0
+            in_dim = padding + in_dim + padding
         if isinstance(in_dim, Dim):
             filter_left_dilated = (filter_size - 1) * dilation_rate // 2
             filter_right_dilated = (filter_size - 1) * dilation_rate - filter_left_dilated
@@ -726,4 +789,95 @@ def _calc_out_dim(in_dim, filter_size, stride, padding, dilation_rate=1):
             return valid_part.ceildiv_right(stride)
         return ceildiv(in_dim - (filter_size - 1) * dilation_rate, stride)
     else:
-        raise Exception("invalid padding %r" % padding)
+        raise ValueError(f"invalid padding {padding!r} (type {type(padding).__name__})")
+
+
+def _should_use_consistent_same_padding() -> bool:
+    """
+    :return: whether to use the new consistent same padding with :func:`_consistent_same_padding`.
+
+    This is only needed for the case when we have striding and padding="same".
+    See :func:`_consistent_same_padding` for more details.
+
+    Check the global RETURNN config for the ``rf_use_consistent_same_padding``
+    on how we should handle the ``padding="same"`` case for convolution/pooling when there is striding.
+    If that is not specified, with behavior version >=24, we will use the new consistent same padding,
+    with behavior version <=23, we will not use it.
+
+    See issue `#1693 <https://github.com/rwth-i6/returnn/issues/1693>`__.
+    """
+    from returnn.config import get_global_config
+
+    config = get_global_config(raise_exception=False)
+    config_value = None
+    if config:
+        if "rf_use_consistent_same_padding" in config.typed_dict:
+            config_value = config.typed_dict["rf_use_consistent_same_padding"]
+            assert config_value is None or isinstance(config_value, bool)
+        elif "rf_use_consistent_same_padding" in config.dict:
+            config_value = config.bool("rf_use_consistent_same_padding", None)
+    if config_value is not None:
+        return config_value
+
+    from returnn.util.basic import BehaviorVersion
+
+    return BehaviorVersion.get() >= 24
+
+
+def _consistent_same_padding(
+    source: Tensor,
+    *,
+    in_spatial_dims: Sequence[Dim],
+    filter_size: Optional[Union[int, Dim, Sequence[int], Sequence[Dim]]],
+    dilation_rate: Optional[Union[int, Sequence[int]]] = None,
+    pad_value: Union[int, float],
+) -> Tuple[Tensor, Sequence[Dim], Union[int, Sequence[int]]]:
+    """
+    In case of striding and padding="same", the standard padding that we do (following TensorFlow)
+    depends on the current dimension size.
+    It adds padding left and right such that the first and last window
+    will have the same amount of padding (+-1).
+    With stride=1, this is the standard (filter_size-1)/2 left and right padding,
+    but with stride>1, this is not the case anymore.
+    (See also the explanation and calculation of padding in :func:`returnn.torch.frontend._backend.TorchBackend.conv`.)
+    However, the problem with this behavior is with batching:
+    The padding now depends on the longest sequence in the batch,
+    and thus is arbitrary for any of the other sequences.
+
+    The new consistent same padding adds padding independent of the current dimension size (largest seq in batch).
+    We just do the same as with stride=1, i.e. (filter_size-1)/2 left and right padding.
+
+    :return: source or padded source, in_spatial_dims or new in_spatial_dims, new padding on top of the output
+    """
+    filter_size = _make_sequence(filter_size or 1, nd=len(in_spatial_dims))
+    dilation_rate = _make_sequence(dilation_rate or 1, nd=len(in_spatial_dims))
+    filter_size_ints = [s.dimension if isinstance(s, Dim) else s for s in filter_size]
+    if all(s % 2 == 1 for s in filter_size_ints):
+        # In this case, we can pass padding as integer to the backend, so that it adds the same padding left/right.
+        return source, in_spatial_dims, [(s // 2) * d for s, d in zip(filter_size_ints, dilation_rate)]
+    # Need to use the custom padding here.
+    paddings = []
+    for s, d in zip(filter_size, dilation_rate):
+        pad_left = (s - 1) * d // 2
+        pad_right = (s - 1) * d - pad_left
+        paddings.append((pad_left, pad_right))
+    # We expect that masking was already done before (or we don't care about it), thus handle_dynamic_dims=False.
+    source, in_spatial_dims = rf.pad(
+        source, axes=in_spatial_dims, padding=paddings, value=pad_value, handle_dynamic_dims=False
+    )
+    return source, in_spatial_dims, 0
+
+
+def _make_sequence(value: Union[int, Sequence[int]], *, nd: int) -> Sequence[int]:
+    if isinstance(value, int):
+        return [value] * nd
+    assert len(value) == nd
+    return value
+
+
+def _any_is_non_default(single_or_seq: Optional[Union[int, Sequence[int]]], *, default: int) -> bool:
+    if single_or_seq is None:
+        return False
+    if isinstance(single_or_seq, int):
+        return single_or_seq != default
+    return any(i != default for i in single_or_seq)

returnn/tf/frontend_layers/_backend.py

@@ -998,7 +998,7 @@ class ReturnnLayersBackend(Backend[Layer]):
         out_spatial_dims: Optional[Sequence[Dim]] = None,
         filter: Tensor,
         filter_size: Sequence[Dim],  # to have the order well-defined
-        padding: str,
+        padding: Union[str, int, Sequence[int]],
         strides: Optional[Union[int, Sequence[int]]] = None,
         dilation_rate: Optional[Union[int, Sequence[int]]] = None,
         groups: Optional[int] = None,
@@ -1088,7 +1088,7 @@ class ReturnnLayersBackend(Backend[Layer]):
         *,
         mode: str,
         pool_size: Sequence[int],
-        padding: str = "valid",
+        padding: Union[str, int, Sequence[int]] = "valid",
         dilation_rate: Union[Sequence[int], int] = 1,
         strides: Sequence[int],
         in_spatial_dims: Sequence[Dim],

returnn/tf/layers/basic.py

@@ -4184,7 +4184,9 @@ class PadLayer(_ConcatInputLayer):
         self,
         *,
         axes: Union[Dim, str, Sequence[Union[Dim, str]]],
-        padding: Union[int, Tuple[int, int], Sequence[Tuple[int, int]]],
+        padding: Union[
+            int, Dim, Tuple[Union[int, Dim], Union[int, Dim]], Sequence[Tuple[Union[int, Dim], Union[int, Dim]]]
+        ],
         out_dims: Optional[Union[Dim, Sequence[Dim]]] = None,
         handle_dynamic_dims: Optional[bool] = None,
         value: Union[int, float] = 0,
@@ -4211,7 +4213,10 @@ class PadLayer(_ConcatInputLayer):
         padding = self._transform_padding(padding=padding, axes=axes)
         paddings = [(0, 0)] * len(range(self.input_data.batch_ndim))
         for i, a in enumerate(axes):
-            paddings[a] = padding[i]
+            pad_left, pad_right = padding[i]
+            pad_left = pad_left.dimension if isinstance(pad_left, Dim) else pad_left
+            pad_right = pad_right.dimension if isinstance(pad_right, Dim) else pad_right
+            paddings[a] = (pad_left, pad_right)
         mode = mode.lower()
         if handle_dynamic_dims is None:
             handle_dynamic_dims = self._handle_dynamic_dims_default(
@@ -4219,7 +4224,7 @@ class PadLayer(_ConcatInputLayer):
                 padding=padding,
                 mode=mode,
             )
-        if all(sum(p) == 0 for p in padding):
+        if all(left == right == 0 for left, right in paddings):
             self.output.placeholder = self.input_data.placeholder
         elif mode == "replication":
             self.output.placeholder = tf_util.pad_replicate(self.input_data.placeholder, axes, padding)
@@ -4227,7 +4232,7 @@ class PadLayer(_ConcatInputLayer):
             self.output.placeholder = tf.pad(
                 self.input_data.placeholder, paddings=paddings, mode=mode, constant_values=value
             )
-        if all(right == 0 for left, right in padding) and mode != "circular":
+        if all(right == 0 for left, right in paddings) and mode != "circular":
             pass  # no masking needed
         else:
             import returnn.frontend as rf
@@ -4257,9 +4262,9 @@ class PadLayer(_ConcatInputLayer):
     @classmethod
     def _transform_padding(cls, padding, axes):
         """
-        :param list[(int,int)]|(int,int)|int padding:
+        :param Sequence[(int|Dim,int|Dim)]|(int|Dim,int|Dim)|int|Dim padding:
         :param list[int] axes:
-        :rtype: list[(int,int)]
+        :rtype: Sequence[(int|Dim,int|Dim)]
         """
         if isinstance(padding, (list, tuple)):
             if isinstance(padding[0], (list, tuple)):
@@ -4316,9 +4321,9 @@ class PadLayer(_ConcatInputLayer):
         """
         :param str name:
         :param list[LayerBase] sources:
-        :param Dim|str|list[Dim|str] axes:
-        :param list[(int,int)]|(int,int)|int padding:
-        :param Dim|list[Dim]|None out_dims:
+        :param Dim|str|Sequence[Dim|str] axes:
+        :param Sequence[(int|Dim,int|Dim)]|(int|Dim,int|Dim)|int|Dim padding:
+        :param Dim|Sequence[Dim]|None out_dims:
         :rtype: Data
         """
         from ..util.data import Dim
@@ -6223,7 +6228,7 @@ class ConvLayer(_ConcatInputLayer):
             for 1D/2D/3D conv.
             The input data ndim must match, or you can add dimensions via input_expand_dims or input_add_feature_dim.
             It will automatically swap the batch-dim to the first axis of the input data.
-        :param str padding: "same", "valid" or "same_static".
+        :param str|int|Sequence[int] padding: "same", "valid" or "same_static".
             "same_static" is calculated differently depending on whether an axis is static or dynamic.
             For static axes, "same_static" padding is the same as "same" padding,
             i.e. filter_size - 1 - (T + strides - 1) % strides.
@@ -6261,8 +6266,10 @@ class ConvLayer(_ConcatInputLayer):
         """
         from returnn.util import BehaviorVersion
 
-        padding = padding.upper()
-        assert padding in ["SAME", "VALID", "SAME_STATIC"], "no other padding supported at the moment"
+        padding = padding.upper() if isinstance(padding, str) else padding
+        assert padding in ["SAME", "VALID", "SAME_STATIC"] or isinstance(
+            padding, (int, tuple, list)
+        ), f"{self}: got unsupported padding {padding}"
         assert "out_type" not in kwargs, "don't set out_type explicitly for this layer"
         assert len(filter_size) in (1, 2, 3), "only 1D conv, 2D conv or 3D conv supported"
         super(ConvLayer, self).__init__(in_dim=in_dim, out_dim=out_dim, **kwargs)
@@ -6398,6 +6405,17 @@ class ConvLayer(_ConcatInputLayer):
                 out_batch_feature_major=out_batch_feature_major,
             )
             padding = "VALID"  # input is now already "same" padded, therefore use "valid" padding from here
+        elif isinstance(padding, int) and padding == 0:
+            x = input_data.placeholder
+            padding = "VALID"
+        elif isinstance(padding, (int, list, tuple)):
+            x = self.get_input_placeholder_with_int_padding(
+                input_data=input_data,
+                num_batch_dims=num_batch_dims,
+                out_batch_feature_major=out_batch_feature_major,
+                padding=padding,
+            )
+            padding = "VALID"
         else:
             x = input_data.placeholder
 
@@ -6539,7 +6557,7 @@ class ConvLayer(_ConcatInputLayer):
         :param Sequence[int|Dim] filter_size:
         :param Sequence[int] strides:
         :param Sequence[int] dilation_rate:
-        :param str padding:
+        :param str|int|Sequence[int] padding:
         """
         if output.feature_dim_axis == num_batch_dims:
             out_spatial_dims_ = output.dim_tags[num_batch_dims + 1 :]
@@ -6558,7 +6576,7 @@ class ConvLayer(_ConcatInputLayer):
                 filter_size=filter_size[i],
                 stride=strides[i],
                 dilation_rate=dilation_rate[i],
-                padding=padding,
+                padding=padding if isinstance(padding, (str, int)) else padding[i],
             )
             assert isinstance(out_tag_calc, Dim)
             out_tag_calc.declare_same_as(out_tag)
@@ -6717,7 +6735,7 @@ class ConvLayer(_ConcatInputLayer):
         """
         Returns the placeholder of input_data with same_static padding applied to it.
 
-        :param input_data:
+        :param input_data: [Batch..., Spatial..., Feature] or [Batch..., Feature, Spatial...]
         :param num_batch_dims:
         :param filter_size:
         :param strides:
@@ -6757,6 +6775,44 @@ class ConvLayer(_ConcatInputLayer):
         x = tf.pad(input_data.placeholder, paddings)
         return x
 
+    @classmethod
+    def get_input_placeholder_with_int_padding(
+        cls,
+        input_data: Data,
+        *,
+        num_batch_dims: int,
+        out_batch_feature_major: bool,
+        padding: Union[int, Sequence[int]],
+        pad_value: float = 0.0,
+    ) -> tf.Tensor:
+        """
+        Returns the placeholder of input_data with same_static padding applied to it.
+
+        :param input_data: [Batch..., Spatial..., Feature] or [Batch..., Feature, Spatial...]
+        :param num_batch_dims:
+        :param out_batch_feature_major:
+        :param padding:
+        :param pad_value:
+        """
+        num_spatial_dims = input_data.batch_ndim - num_batch_dims - 1
+        if isinstance(padding, int):
+            padding = [padding] * num_spatial_dims
+        paddings = [[0, 0] for _ in range(input_data.batch_ndim)]
+        for axis, dim in enumerate(input_data.dims):
+            if axis < num_batch_dims:
+                continue
+            if axis == num_batch_dims and out_batch_feature_major:
+                # input_data has dimensions [batch] * num_batch_dims + [channels] + [spatial] * num_spatial_dims
+                continue
+            if axis >= num_batch_dims + num_spatial_dims and not out_batch_feature_major:
+                # input_data has dimensions [batch] * num_batch_dims + [spatial] * num_spatial_dims + [channels]
+                break
+
+            padding_ = padding[axis - num_batch_dims - out_batch_feature_major]
+            paddings[axis] = [padding_, padding_]
+        x = tf.pad(input_data.placeholder, paddings, constant_values=pad_value)
+        return x
+
     @classmethod
     def calc_out_dim(cls, in_dim, filter_size, stride, padding, dilation_rate=1):
         """
@@ -6764,7 +6820,7 @@ class ConvLayer(_ConcatInputLayer):
         :param int|Dim filter_size: e.g. 2, for the corresponding axis
         :param int stride: e.g. 1, for the corresponding axis
         :param int dilation_rate: e.g. 1
-        :param str padding: "valid" or "same"
+        :param str|int padding: "valid" or "same"
         :return: the output dimension
         :rtype: T
         """
@@ -6779,13 +6835,16 @@ class ConvLayer(_ConcatInputLayer):
                 return a
             return -(-a // b)
 
-        padding = padding.upper()
+        padding = padding.upper() if isinstance(padding, str) else padding
         # See tf.compat.v1.nn.convolution() documentation for more.
         if padding == "SAME":
             if isinstance(in_dim, Dim):
                 return in_dim.ceildiv_right(stride)
             return ceildiv(in_dim, stride)
-        elif padding == "VALID":
+        elif padding == "VALID" or isinstance(padding, int):
+            if isinstance(padding, int) and padding != 0:
+                assert padding > 0
+                in_dim = padding + in_dim + padding
             if isinstance(in_dim, Dim):
                 filter_left_dilated = (filter_size - 1) * dilation_rate // 2
                 filter_right_dilated = (filter_size - 1) * dilation_rate - filter_left_dilated
@@ -6826,7 +6885,7 @@ class ConvLayer(_ConcatInputLayer):
         :param Sequence[LayerBase] sources:
         :param returnn.tf.network.TFNetwork network:
         :param Sequence[int|Dim] filter_size:
-        :param str padding:
+        :param str|int|Sequence[int] padding:
         :param int|Sequence[int] strides:
         :param int|Sequence[int] dilation_rate:
         :param int input_expand_dims: number of dynamic dims to add to the input
@@ -6839,6 +6898,7 @@ class ConvLayer(_ConcatInputLayer):
         :param Sequence[Dim]|None out_spatial_dims:
         :param int input_expand_dims: number of spatial dims to add to the input
         :param bool|NotSpecified auto_use_channel_first:
+        :rtype: Data
         """
         from returnn.util import BehaviorVersion
 
@@ -6857,7 +6917,8 @@ class ConvLayer(_ConcatInputLayer):
         assert len(dilation_rate) == len(filter_size)
         if in_spatial_dims:
             assert len(in_spatial_dims) == len(filter_size)
-        padding = padding.upper()
+        if isinstance(padding, str):
+            padding = padding.upper()
         # Be relaxed about incorrect input data. Throw errors later. This can also work during template construction.
         if not input_data.have_batch_axis():
             input_data = input_data.copy_add_batch_dim(batch_dim_axis=0)
@@ -6889,7 +6950,11 @@ class ConvLayer(_ConcatInputLayer):
             for i in range(len(filter_size)):
                 old_tag = old_spatial_dim_tags[i] if i < len(old_spatial_dim_tags) else None
                 filter_size_ = filter_size[i].dimension if isinstance(filter_size[i], Dim) else filter_size[i]
-                if old_tag and (filter_size_ == strides[i] == 1 or (strides[i] == 1 and padding == "SAME")):
+                padding_ = padding if isinstance(padding, (str, int)) else padding[i]
+                if old_tag and (
+                    (filter_size_ == strides[i] == 1 and padding_ in ("SAME", "VALID", 0))
+                    or (strides[i] == 1 and padding_ == "SAME")
+                ):
                     dim_tags.append(old_tag)  # identity in this axis
                     continue
                 new_dim = None
@@ -6899,7 +6964,7 @@ class ConvLayer(_ConcatInputLayer):
                         filter_size=filter_size[i],
                         stride=strides[i],
                         dilation_rate=dilation_rate[i],
-                        padding=padding,
+                        padding=padding_,
                     )
                 dim_tags.append(
                     Dim(
@@ -7009,8 +7074,8 @@ class PoolLayer(_ConcatInputLayer):
     ):
         """
         :param str mode: "max" or "avg"
-        :param tuple[int] pool_size: shape of the window of each reduce
-        :param str padding: "same", "valid" or "same_static".
+        :param Sequence[int] pool_size: shape of the window of each reduce
+        :param str|int|Sequence[int] padding: "same", "valid" or "same_static".
             "same_static" is calculated differently depending on whether an axis is static or dynamic.
             For static axes, "same_static" padding is the same as "same" padding,
             i.e. filter_size - 1 - (T + strides - 1) % strides.
@@ -7018,13 +7083,13 @@ class PoolLayer(_ConcatInputLayer):
             filter_size - 1, i.e. it is independent of the length T of the axis and the striding.
             For dynamic axes, to avoid skipping any frames on the right,
             we set left_padding = (filter_size - strides) // 2.
-        :param tuple[int]|int dilation_rate:
-        :param tuple[int]|int|None strides: in contrast to tf.nn.pool, the default (if it is None)
+        :param Sequence[int]|int dilation_rate:
+        :param Sequence[int]|int|None strides: in contrast to tf.nn.pool, the default (if it is None)
             will be set to pool_size
         :param Dim|None in_dim:
-        :param list[Dim|str]|None in_spatial_dims:
+        :param Sequence[Dim|str]|None in_spatial_dims:
         :param Dim|None out_dim:
-        :param list[Dim]|None out_spatial_dims:
+        :param Sequence[Dim]|None out_spatial_dims:
         :param bool|NotSpecified use_channel_first: if set, will transform input to NCHW format
         :param bool use_time_mask:
         """
@@ -7032,8 +7097,15 @@ class PoolLayer(_ConcatInputLayer):
         assert "out_type" not in kwargs
         mode = mode.upper()
         assert mode in ["MAX", "AVG"]
-        padding = padding.upper()
-        assert padding in ["VALID", "SAME", "SAME_STATIC"]
+        if isinstance(padding, str):
+            padding = padding.upper()
+            assert padding in ["VALID", "SAME", "SAME_STATIC"]
+        elif isinstance(padding, int) or (
+            isinstance(padding, (list, tuple)) and all(isinstance(p, int) for p in padding)
+        ):
+            pass
+        else:
+            raise TypeError(f"invalid type ({type(padding).__name__}) for padding: {padding}")
         if isinstance(dilation_rate, int):
             dilation_rate = [dilation_rate] * len(pool_size)
         assert len(dilation_rate) == len(pool_size)
@@ -7102,6 +7174,18 @@ class PoolLayer(_ConcatInputLayer):
                 out_batch_feature_major=out_batch_feature_major,
             )
             padding = "VALID"  # input is now already "same" padded, therefore use "valid" padding from here
+        elif isinstance(padding, int) and padding == 0:
+            x = input_data.placeholder
+            padding = "VALID"
+        elif isinstance(padding, (int, list, tuple)):
+            x = ConvLayer.get_input_placeholder_with_int_padding(
+                input_data=input_data,
+                num_batch_dims=num_batch_dims,
+                out_batch_feature_major=out_batch_feature_major,
+                padding=padding,
+                pad_value={"MAX": float("-inf"), "AVG": 0}[mode],
+            )
+            padding = "VALID"
         else:
             x = input_data.placeholder
 
@@ -7145,14 +7229,14 @@ class PoolLayer(_ConcatInputLayer):
         :param str name:
         :param list[LayerBase] sources:
         :param returnn.tf.network.TFNetwork network:
-        :param tuple[int]|list[int] pool_size:
-        :param tuple[int]|list[int]|int strides:
-        :param int|tuple[int]|list[int] dilation_rate:
-        :param str padding:
+        :param Sequence[int] pool_size:
+        :param Sequence[int]|int strides:
+        :param int|Sequence[int] dilation_rate:
+        :param str|int|Sequence[int] padding:
         :param Dim|None in_dim:
-        :param list[Dim|str]|None in_spatial_dims:
+        :param Sequence[Dim|str]|None in_spatial_dims:
         :param Dim|None out_dim:
-        :param list[Dim]|None out_spatial_dims:
+        :param Sequence[Dim]|None out_spatial_dims:
         :param bool|NotSpecified use_channel_first:
         :rtype: Data
         """

returnn/tf/util/basic.py

@@ -3733,13 +3733,14 @@ def single_strided_slice(x, axis, begin=None, end=None, step=None):
 def pad_replicate(x, axes, padding):
     """
     :param tf.Tensor x:
-    :param list[int] axes:
-    :param list[(int,int)] padding:
+    :param Sequence[int] axes:
+    :param Sequence[(int|Dim,int|Dim)] padding:
     :rtype: tf.Tensor
     """
     with tf.name_scope("pad_replicate"):
         assert len(padding) == 1, "Not implemented otherwise yet"
         assert len(axes) == 1, "Not implemented otherwise yet"
+        assert isinstance(padding[0][0], int) and isinstance(padding[0][1], int)  # not implemented otherwise yet
         pad_left = tf.gather(x, 0, axis=axes[0])
         pad_left = tf.expand_dims(pad_left, axis=axes[0])
         pad_left = tf.repeat(pad_left, padding[0][0], axis=axes[0])

returnn/torch/engine.py

@@ -3,7 +3,7 @@ Main engine for PyTorch
 """
 
 from __future__ import annotations
-from typing import Optional, Any, Union, Callable, Dict, Set, Tuple
+from typing import Optional, Any, Union, Callable, Dict, Set
 from contextlib import nullcontext, ExitStack, contextmanager
 
 import gc
@@ -365,8 +365,6 @@ class Engine(EngineBase):
         zero_grad_next_step = True
         cur_count_grad_accum = 0
         extern_data = None
-        num_seqs = None
-        last_seq_idx = 0
 
         total_data_size_packed = NumbersDict()
         total_data_size_padded = NumbersDict()
@@ -400,20 +398,8 @@ class Engine(EngineBase):
                 )
 
                 complete_frac = float(extern_data_raw["complete_frac"])
-                num_seqs, last_seq_idx = _get_num_seqs_last_seq_idx(
-                    report_prefix=report_prefix,
-                    extern_data_raw=extern_data_raw,
-                    step_idx=step_idx,
-                    prev_num_seqs=num_seqs,
-                    prev_last_seq_idx=last_seq_idx,
-                )
-                epoch_continuous = (
-                    self.epoch - 1 + complete_frac
-                    if complete_frac >= 0.0
-                    else (self.epoch - 1 + (last_seq_idx + 1) / num_seqs)
-                    if num_seqs is not None
-                    else None
-                )
+                epoch_continuous = self.epoch - 1 + complete_frac if complete_frac >= 0.0 else None
+                num_seqs = int(extern_data_raw["num_seqs"])
 
                 # clear the gradients when every gradient accumulation loop starts
                 if zero_grad_next_step:
@@ -490,7 +476,7 @@ class Engine(EngineBase):
                     eval_info=dict(eval_info),
                     step_duration=step_duration,
                     start_elapsed=step_end_time - epoch_start_time,
-                    seq_idx=last_seq_idx,
+                    complete_frac=complete_frac,
                     num_seqs=num_seqs,
                     batch_size_info=_get_batch_size_info(extern_data) if self._log_batch_size else None,
                     log_memory_usage_device=self._device if self._log_memory_usage else None,
@@ -629,13 +615,18 @@ class Engine(EngineBase):
             accumulated_losses_dict = NumbersDict()
             accumulated_inv_norm_factors_dict = NumbersDict()
             step_idx = 0
+            eval_start_time = time.monotonic()
 
+            report_prefix = f"ep {self.epoch} {dataset_name} eval"
             with torch.no_grad():
                 for extern_data_raw in data_loader:
                     if self._torch_distributed_ctx and step_idx % 100 == 0:
                         _has_data = torch.tensor([True], device="cpu", dtype=torch.int8)
                         torch.distributed.broadcast(_has_data, src=0)
 
+                    complete_frac = float(extern_data_raw["complete_frac"])
+                    num_seqs = int(extern_data_raw["num_seqs"])
+
                     extern_data = extern_data_util.raw_dict_to_extern_data(
                         extern_data_raw,
                         extern_data_template=self.extern_data,
@@ -644,6 +635,8 @@ class Engine(EngineBase):
                     )
 
                     self._run_step(extern_data, train_func=True)
+                    step_end_time = time.monotonic()
+
                     train_ctx = rf.get_run_ctx()
 
                     losses_dict = NumbersDict(
@@ -664,9 +657,12 @@ class Engine(EngineBase):
                     accumulated_inv_norm_factors_dict += inv_norm_factors_dict
                     eval_info = self._maybe_extend_losses_info(losses_dict / inv_norm_factors_dict)
                     _print_process(
-                        f"ep {self.epoch} {dataset_name} eval",
+                        report_prefix,
                         step=step_idx,
                         eval_info=dict(eval_info),
+                        complete_frac=complete_frac,
+                        num_seqs=num_seqs,
+                        start_elapsed=step_end_time - eval_start_time,
                         log_memory_usage_device=self._device if self._log_memory_usage else None,
                     )
                     step_idx += 1
@@ -1290,8 +1286,6 @@ class Engine(EngineBase):
             new_dim.dyn_size_ext = _get_tensor_wo_batch_numpy(dim.dyn_size_ext)
             return new_dim
 
-        num_seqs = None
-        last_seq_idx = 0
         report_prefix = f"ep {self.epoch} {dataset.name} forward"
         with torch.no_grad():
             callback.init(model=self._orig_model)
@@ -1300,13 +1294,8 @@ class Engine(EngineBase):
             for extern_data_raw in data_loader:
                 step_begin_time = time.monotonic()
 
-                num_seqs, last_seq_idx = _get_num_seqs_last_seq_idx(
-                    report_prefix=report_prefix,
-                    extern_data_raw=extern_data_raw,
-                    step_idx=step_idx,
-                    prev_num_seqs=num_seqs,
-                    prev_last_seq_idx=last_seq_idx,
-                )
+                complete_frac = float(extern_data_raw["complete_frac"])
+                num_seqs = int(extern_data_raw["num_seqs"])
 
                 if self._forward_step_expected_outputs:
                     # Also resets any dyn dims, which might have been set in the prev step.
@@ -1354,7 +1343,7 @@ class Engine(EngineBase):
                     eval_info=None,
                     step_duration=step_duration,
                     start_elapsed=step_end_time - epoch_start_time,
-                    seq_idx=last_seq_idx,
+                    complete_frac=complete_frac,
                     num_seqs=num_seqs,
                     batch_size_info=_get_batch_size_info(extern_data) if self._log_batch_size else None,
                     log_memory_usage_device=self._device if self._log_memory_usage else None,
@@ -1442,7 +1431,7 @@ def _print_process(
     batch_size_info: Optional[Dict[str, Any]] = None,
     step_duration: Optional[float] = None,
     start_elapsed: Optional[float] = None,
-    seq_idx: Optional[int] = None,
+    complete_frac: Optional[float] = None,
     num_seqs: Optional[int] = None,
     log_memory_usage_device: Optional[str] = None,
 ):
@@ -1455,11 +1444,14 @@ def _print_process(
     :param batch_size_info:
     :param step_duration: time elapsed for this step (secs)
     :param start_elapsed: time elapsed since epoch start (secs)
-    :param num_seqs: total number of sequences for this epoch
+    :param complete_frac: how much of the current epoch is already consumed
+    :param num_seqs: total number of seqs this epoch
     :param log_memory_usage_device: if given, will log memory usage (peak allocated memory)
     :return: nothing, will be printed to log
     """
     if log.verbose[5]:  # report every minibatch
+        if step == 0 and num_seqs is not None and num_seqs >= 0:
+            print(f"{report_prefix} num_seqs: {num_seqs}", file=log.v5)
         info = [report_prefix, "step %i" % step]
         if eval_info:  # Such as score.
             info += ["%s %s" % (k, _format_score_value(v)) for k, v in eval_info.items()]
@@ -1475,17 +1467,16 @@ def _print_process(
             info += ["%.3f sec/step" % step_duration]
         if start_elapsed is not None:
             info += ["elapsed %s" % hms(start_elapsed)]
-        if num_seqs is not None:
-            assert seq_idx is not None and start_elapsed is not None  # unexpected combination...
-            complete = (seq_idx + 1) / num_seqs
-            assert 1 >= complete > 0, f"{step} step, {num_seqs} num_seqs"
-            total_time_estimated = start_elapsed / complete
+        if complete_frac is not None:
+            assert 1 >= complete_frac > 0, f"{step} step, {complete_frac} complete_frac"
+            assert start_elapsed is not None
+            total_time_estimated = start_elapsed / complete_frac
             remaining_estimated = total_time_estimated - start_elapsed
             info += [
                 "exp. remaining %s" % hms(remaining_estimated),
-                "complete %.02f%%" % (complete * 100),
+                "complete %.02f%%" % (complete_frac * 100),
             ]
-        if start_elapsed is not None and num_seqs is None:
+        if start_elapsed is not None and complete_frac is None:
             info += ["(unk epoch len)"]
         print(", ".join(filter(None, info)), file=log.v5)
 
@@ -1634,27 +1625,3 @@ def _get_total_grad_norm(model: torch.nn.Module, p: float) -> float:
             p=p,
         ).item()
     )
-
-
-def _get_num_seqs_last_seq_idx(
-    *,
-    report_prefix: str,
-    extern_data_raw: Dict[str, Any],
-    step_idx: int,
-    prev_num_seqs: Optional[int],
-    prev_last_seq_idx: int,
-) -> Tuple[Optional[int], int]:
-    num_seqs = prev_num_seqs
-    num_seqs_ = int(extern_data_raw["num_seqs"]) if extern_data_raw.get("num_seqs", None) is not None else -1
-    # Note: The batches might have been shuffled,
-    # thus we cannot really assert that the seq_idx is always increasing.
-    last_seq_idx = max(int(extern_data_raw["seq_idx"].max()), prev_last_seq_idx)
-    if step_idx == 0:
-        if num_seqs_ >= 0:
-            print(f"{report_prefix} num_seqs: {num_seqs_}", file=log.v5)
-            num_seqs = num_seqs_
-    elif num_seqs_ >= 0:
-        assert num_seqs_ == num_seqs
-    if num_seqs is not None:
-        assert last_seq_idx < num_seqs
-    return num_seqs, last_seq_idx

returnn/torch/frontend/_backend.py

@@ -1879,7 +1879,7 @@ class TorchBackend(Backend[torch.Tensor]):
         out_spatial_dims: Optional[Sequence[Dim]] = None,
         filter: Tensor,
         filter_size: Sequence[Dim],  # to have the order well-defined
-        padding: str,
+        padding: Union[str, int, Sequence[int]],
         strides: Optional[Union[int, Sequence[int]]] = None,
         dilation_rate: Optional[Union[int, Sequence[int]]] = None,
         groups: Optional[int] = None,
@@ -2008,7 +2008,7 @@ class TorchBackend(Backend[torch.Tensor]):
         *,
         mode: str,
         pool_size: Sequence[int],
-        padding: str = "valid",
+        padding: Union[str, int, Sequence[int]] = "valid",
         dilation_rate: Union[Sequence[int], int] = 1,
         strides: Sequence[int],
         in_spatial_dims: Sequence[Dim],
@@ -2035,19 +2035,22 @@ class TorchBackend(Backend[torch.Tensor]):
             [-1, batch_dims[-1].get_dim_value() if batch_dims else 1] + [d.get_dim_value() for d in in_spatial_dims],
         )
         assert isinstance(strides, (list, tuple)) and len(strides) == len(in_spatial_dims) == len(pool_size)
-        if padding.lower() == "same":
+        if isinstance(padding, str) and padding.lower() == "same":
             # padding='same' is not quite the same as ceil_mode=True, so we explicitly pad here.
             padding = []
             for i, s in enumerate(pool_size):
                 # See comment in conv.
+                # I'm a bit unsure here... https://github.com/pytorch/pytorch/issues/148123
                 pad = s - 1 - (src_raw.shape[2 + i] - 1) % strides[i]
                 padding.append(pad // 2)
             ceil_mode = True
-        elif padding.lower() == "valid":
+        elif isinstance(padding, str) and padding.lower() == "valid":
             padding = 0
             ceil_mode = False
+        elif isinstance(padding, (int, tuple, list)):
+            ceil_mode = False
         else:
-            raise ValueError(f"invalid padding {padding!r}")
+            raise ValueError(f"invalid padding {padding!r} (type {type(padding).__name__}")
         func_name = f"{mode}_pool{len(in_spatial_dims)}d"
         func = getattr(torch.nn.functional, func_name)  # e.g. torch.nn.functional.max_pool1d
         kwargs = {}

returnn/util/basic.py

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

{returnn-1.20250304.10039.dist-info → returnn-1.20250304.113330.dist-info}/METADATA

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

{returnn-1.20250304.10039.dist-info → returnn-1.20250304.113330.dist-info}/RECORD

@@ -1,9 +1,9 @@
-returnn/PKG-INFO,sha256=et7Z9NstTVvnWjiIMXhquw3eiMnMxYMfnEEbVc755xQ,5214
+returnn/PKG-INFO,sha256=BmSxZKkRxyL20E4Zsud1muiQ-rth9Ob9PMR-43IrAMw,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=d4hd9PkngTUKLJT4Q6GLMhVg4nXyV3Pym04_IKcblgc,77
+returnn/_setup_info_generated.py,sha256=94BElbYUGmjpsoY8BzvfW39RUTXw9Fy3UwlPoEjrkU8,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
@@ -75,7 +75,7 @@ returnn/extern/graph_editor/subgraph.py,sha256=R3uIFqWgiL7L5S4YATm9o9a3wfEa_mSb4
 returnn/extern/graph_editor/transform.py,sha256=d9fEgu0JC342q0g9niVxRWMKzkQQA9mrrajBGcU1o_s,29349
 returnn/extern/graph_editor/util.py,sha256=QMrQeQZ7lJwsrNQub9tof0h3quEaoHiGJaZmogQ7jXE,18707
 returnn/frontend/__init__.py,sha256=2aS7nbxXniIrBp2DODl0xN0f3IJ_dX4Bi9ZlR7W5_DE,1472
-returnn/frontend/_backend.py,sha256=TNkEdj9GKxJfSM1ZMQ_SdAQzn2TU7SQbG6JGdaWhUeI,50374
+returnn/frontend/_backend.py,sha256=JNqQomHCN4-1VLq5o9VRbs_L8gSZkvOgjUmRYt8jx1o,50428
 returnn/frontend/_cache.py,sha256=JAhi7L-raQ3A-NC3JUYDtdRTwT3BGJJGGZxrZ8MfEWQ,8403
 returnn/frontend/_numpy_backend.py,sha256=fZjks7p3dgxVZ6tSDazTTgBxNjJqXjfqgw_7mA7rDEE,9066
 returnn/frontend/_random_journal.py,sha256=_ktP_mjgx8vtQQGX_DofdhewJj0aPiczefTWeemPkmo,5457
@@ -88,7 +88,7 @@ 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=p4R6j40GCvVrw3kbQQJtfxY6tfIR8Rb3tIzwAtiLuec,23858
+returnn/frontend/conv.py,sha256=Q0q90-uu9d6qV-v8_DlFGxpZtc6FjfXVpfkkXmv1Alk,31959
 returnn/frontend/device.py,sha256=K7Y1qoQcO4GIHgLkPLQWK-GVT8gKL8GwyQrmPo8LgBE,1438
 returnn/frontend/dims.py,sha256=aH5FQ_m0xMD6Rj-BUWGx8lB-HkCuwZfMBf6mZbGGW5E,12611
 returnn/frontend/dropout.py,sha256=rsx3p5b0NblBfXXSQZTQFJ8jUUS3fj4Qzc39iffBMCA,5006
@@ -177,7 +177,7 @@ returnn/tf/sprint.py,sha256=Yqjh0-6sCWHpdDPQCzHKx7TwQCOjJyjfd0KHtnYdd-8,5471
 returnn/tf/updater.py,sha256=St4Z5iBjlkWaB6CiS-K1VNc_iLaan2e6-mVMTTPldzk,72034
 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=U7rbRY9XgMkxxyWY2D8KG-KesSOEGLCxn-Gl6dgwmPc,47277
+returnn/tf/frontend_layers/_backend.py,sha256=igo147YCTVdNuUBm2euEwjAhpH5yDHyQAf5T4jcCrLM,47331
 returnn/tf/frontend_layers/_utils.py,sha256=ijByaDOqPDod5mZC9EoTkt8PHBEODXHsWbkwDOF9XW4,4205
 returnn/tf/frontend_layers/cond.py,sha256=yQ2h5W0sgMZndJdrWv2EE9k9yIcspQ1U0HwBSh3hOKE,14830
 returnn/tf/frontend_layers/config_entry_points.py,sha256=t01RWOiaZohzuqPXX-MLV0P5yCOfE0dz-9dZ77_pK4c,5751
@@ -193,13 +193,13 @@ returnn/tf/frontend_low_level/__init__.py,sha256=34469k3KzMUIGowxReOZnbf6WdTjxY7
 returnn/tf/frontend_low_level/_backend.py,sha256=JwwRRIGnElqBC4bTImdB7w3U1u_SJESeZHYLmq86wog,24479
 returnn/tf/layers/__init__.py,sha256=Ngu-X84nWFgz7ndDu88DqoZ-5lUMMTQWH4g7N8pSoCg,72
 returnn/tf/layers/base.py,sha256=KcADpZUxqLkoFpQPMe_l9thRC7rpyBJIZCHITmnOd7M,153169
-returnn/tf/layers/basic.py,sha256=la0EwaHVzAbL6JOXs6QXnYQ74F3R16piYpT55VwVFT4,611063
+returnn/tf/layers/basic.py,sha256=7eefkCNa8aqh96Hl2Tr8b6rqpE0cudgyyQCuQK-QNKU,615168
 returnn/tf/layers/rec.py,sha256=K9vvyDJeDApYQDKabz7PaOTGHeSTloInkecxKTbqeTU,548357
 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=F1-3Huh4mdoLgBCYbh4z5rDFz2meWMwsGQc3B87wOXg,302811
+returnn/tf/util/basic.py,sha256=8c0xEQNcsIvts2ydwZdUvqk4HsTJFyH_xYPQzaZbV6M,302941
 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
@@ -207,7 +207,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=skFyutdVztxgTEk3HHJ8S83qRWbNpkNT8Tj16Ic0_hE,6981
-returnn/torch/engine.py,sha256=sU9A96icaj65uaEkX4i4aUK3IrB2S19_Fb9_sueB_JE,77426
+returnn/torch/engine.py,sha256=2FLLb2m4sWFwYOQGREDSxQCheCKd_osnFJCdLa_4TzE,76400
 returnn/torch/updater.py,sha256=GqtBvZpElPVMm0lq84JPl4NVLFFETZAzAbR0rTomSao,28249
 returnn/torch/data/__init__.py,sha256=6cLNEi8KoGI12PF6akN7mI_mtjlx-0hcQAfMYoExwik,132
 returnn/torch/data/extern_data.py,sha256=_uT_9_gd5HIh1IoRsrebVG-nufSnb7fgC5jyU05GxJg,7580
@@ -216,7 +216,7 @@ returnn/torch/data/queued_data_iter.py,sha256=PoOsGHdHVZjTmcyfq_ZOw--P6hyfTdmAWI
 returnn/torch/data/returnn_dataset_wrapper.py,sha256=2CaDapzrlqahANuq-nyVAtv5ENHuM8A7okORwYJDisg,8006
 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=rFCoCnzZoBtHPg7mWpO3yJOJMVesuWuA3_6GGSKMc5k,101452
+returnn/torch/frontend/_backend.py,sha256=SKxxpIM0rXEcZ92p-Um5thfC7vmDoZmda13SMAXVYL0,101771
 returnn/torch/frontend/_rand.py,sha256=1JgIkV2XmpgJD86zXZ-NCAe-QuoP2swr6NaS1oz3Qa8,1830
 returnn/torch/frontend/bridge.py,sha256=Z2_UW8AagezC7zsXDc5PKcd8G9WwisV7j9SWGHU0m4U,7840
 returnn/torch/frontend/raw_ops.py,sha256=lF0h-KtYYsdaaqQADylVZp9qzPskOOXA4MfmYDyx5IU,296
@@ -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=Iynt9ATEs_8DaZsX5z6weMyaO2xW9o3gaywq6X7mbEc,142380
+returnn/util/basic.py,sha256=eLlzR-ARGWJoiyRb5-SH5v2zx1jgR_5vuQ5jwYO5Cww,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.20250304.10039.dist-info/LICENSE,sha256=ywBD_U2aD4vpuoIgNAsjIGBYydl0tVKll3De0Z8s77c,11041
-returnn-1.20250304.10039.dist-info/METADATA,sha256=et7Z9NstTVvnWjiIMXhquw3eiMnMxYMfnEEbVc755xQ,5214
-returnn-1.20250304.10039.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-returnn-1.20250304.10039.dist-info/top_level.txt,sha256=Lsn4WZc5Pbfk0-xDQOgnFCxOoqxL4CyeM3N1TFbJncw,8
-returnn-1.20250304.10039.dist-info/RECORD,,
+returnn-1.20250304.113330.dist-info/LICENSE,sha256=ywBD_U2aD4vpuoIgNAsjIGBYydl0tVKll3De0Z8s77c,11041
+returnn-1.20250304.113330.dist-info/METADATA,sha256=BmSxZKkRxyL20E4Zsud1muiQ-rth9Ob9PMR-43IrAMw,5215
+returnn-1.20250304.113330.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+returnn-1.20250304.113330.dist-info/top_level.txt,sha256=Lsn4WZc5Pbfk0-xDQOgnFCxOoqxL4CyeM3N1TFbJncw,8
+returnn-1.20250304.113330.dist-info/RECORD,,