returnn 1.20240611.113715__tar.gz → 1.20240614.213333__tar.gz

{returnn-1.20240611.113715/returnn.egg-info → returnn-1.20240614.213333}/PKG-INFO

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

returnn-1.20240614.213333/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240614.213333'
+long_version = '1.20240614.213333+git.7ddcd01'

{returnn-1.20240611.113715 → returnn-1.20240614.213333}/returnn/frontend/dims.py

@@ -2,7 +2,6 @@
 Utilities for dimension tags, dimensions, axes.
 """
 
-
 from __future__ import annotations
 from typing import Optional, Union, TypeVar, Sequence, Tuple
 from returnn.tensor import Tensor, Dim
@@ -18,6 +17,7 @@ __all__ = [
     "replace_dim",
     "dim_match_priority_when_needed",
     "num_elements_of_shape",
+    "masked_fraction_of_shape",
 ]
 
 
@@ -122,11 +122,19 @@ def dim_match_priority_when_needed(dim: Dim, *other_dims: Dim) -> Dim:
     return dim
 
 
-def num_elements_of_shape(dims: Sequence[Dim]) -> Union[int, Tensor]:
+def num_elements_of_shape(dims: Union[Dim, Sequence[Dim]], *, use_mask: bool = True) -> Union[int, Tensor]:
     """
     :param dims:
+    :param use_mask:
     :return: num elements of a tensor of shape dims, properly considering masking
     """
+    if isinstance(dims, Dim):
+        dims = [dims]
+    if not use_mask:
+        n = 1
+        for dim in dims:
+            n *= dim.get_dim_value_tensor()
+        return n
     if all(dim.is_static() for dim in dims):
         n = 1
         for dim in dims:
@@ -135,7 +143,7 @@ def num_elements_of_shape(dims: Sequence[Dim]) -> Union[int, Tensor]:
 
     n = 1
     dims = list(dims)
-    dims.sort(key=lambda dim: -dim.dyn_size_ext.batch_ndim if dim.dyn_size_ext else 0)
+    dims.sort(key=lambda dim__: -dim__.dyn_size_ext.batch_ndim if dim__.dyn_size_ext else 0)
     while dims:
         dim = dims.pop(0)
         if dim.is_static():
@@ -144,10 +152,29 @@ def num_elements_of_shape(dims: Sequence[Dim]) -> Union[int, Tensor]:
         # E.g. dyn_size_ext is shape [B], and self has shape [B,T].
         # Due to the sorting of dims above, dims will be [T,B], and we will first process T.
         # We want to sum over dyn_size_ext, but then we need to remove the other dims it covers.
+        dims_to_reduce = []
         for dim_ in dim.dyn_size_ext.dims:
-            assert dim_ in dims  # num elements not really well-defined then
-            assert not dim_.need_masking()  # not implemented
-            dims.remove(dim_)
-        n_ = rf.reduce_sum(dim.dyn_size_ext, axis=dim.dyn_size_ext.dims)
+            if dim_ in dims:
+                assert not dim_.need_masking()  # not implemented
+                dims.remove(dim_)
+                dims_to_reduce.append(dim_)
+        n_ = rf.reduce_sum(dim.dyn_size_ext, axis=dims_to_reduce) if dims_to_reduce else dim.dyn_size_ext
         n *= n_
     return n
+
+
+def masked_fraction_of_shape(dims: Union[Dim, Sequence[Dim]], *, inverse: bool = False) -> Union[int, float, Tensor]:
+    """
+    :param dims:
+    :param inverse: if True, return the inverse of the fraction
+    :return: :func:`num_elements_of_shape`(dims) / prod(dims) if not inverse else prod(dims) / num_elements
+    """
+    if isinstance(dims, Dim):
+        dims = [dims]
+    if not any(dim.need_masking() for dim in dims):
+        return 1
+    num_elems_masked = num_elements_of_shape(dims)
+    num_elems_total = 1
+    for dim in dims:
+        num_elems_total *= dim.get_dim_value_tensor()
+    return (num_elems_masked / num_elems_total) if not inverse else (num_elems_total / num_elems_masked)

{returnn-1.20240611.113715 → returnn-1.20240614.213333}/returnn/frontend/normalization.py

@@ -2,7 +2,6 @@
 Normalization functions such as batch norm
 """
 
-
 from __future__ import annotations
 from typing import Optional, Sequence, Union, Tuple
 from returnn.tensor import Tensor, Dim
@@ -12,15 +11,34 @@ import returnn.frontend as rf
 __all__ = ["moments", "LayerNorm", "BatchNorm", "normalize", "Normalize"]
 
 
-def moments(x: Tensor, axis: Union[Dim, Sequence[Dim]]) -> Tuple[Tensor, Tensor]:
+def moments(
+    x: Tensor, axis: Union[Dim, Sequence[Dim]], *, use_mask: bool = True, correction: Union[int, float, Tensor] = 0
+) -> Tuple[Tensor, Tensor]:
     """
     :param x: input
-    :param axis: the axis to be reduced, to calculate statistics over
-    :return: mean, variance. it has the same shape as the input with the axis removed
+    :param axis: the axis (or axes) to be reduced, to calculate statistics over
+    :param use_mask: whether to use a mask for dynamic spatial dims in the reduction
+    :param correction:
+        The variance will be estimated by ``sum((x - mean)**2) / (n-correction)``
+        where ``n`` is the number of elements in the axis (or the axes)
+        (with ``use_mask=True``, taking masking into account, using :func:`num_elements_of_shape`).
+        The default ``correction=0`` will return the biased variance estimation.
+        ``correction=1`` is the `Bessel correction <https://en.wikipedia.org/wiki/Bessel%27s_correction>`__
+        and will return the unbiased variance estimation.
+        In PyTorch, there was an argument ``unbiased`` for this, but this changed recently to ``correction``
+        (`PyTorch issue #61492 <https://github.com/pytorch/pytorch/issues/61492>`__,
+         `Python Array API Standard
+          <https://data-apis.org/array-api/latest/API_specification/generated/array_api.var.html>`__).
+        In PyTorch, the default is ``correction=1``, which is the unbiased variance estimation,
+        while in most other frameworks, the default is ``correction=0``, which is the biased variance estimation.
+    :return: tuple (mean, variance). it has the same shape as the input with the axis removed
     """
     mean = rf.reduce_mean(x, axis=axis)
     # stop_gradient does not change the gradient here
-    variance = rf.reduce_mean(rf.squared_difference(x, rf.stop_gradient(mean)), axis=axis)
+    variance = rf.reduce_mean(rf.squared_difference(x, rf.stop_gradient(mean)), axis=axis, use_mask=use_mask)
+    if isinstance(correction, Tensor) or correction != 0:
+        n = rf.num_elements_of_shape(axis, use_mask=use_mask)
+        variance *= n / (n - correction)
     return mean, variance
 
 

{returnn-1.20240611.113715 → returnn-1.20240614.213333}/returnn/tensor/_tensor_extra.py

@@ -1262,13 +1262,15 @@ class _TensorMixin(_TensorMixinBase):
             raw_tensor = backend.transpose_raw(raw_tensor, [p for p in out_permutation if p >= 0])
             raw_tensor = backend.reshape_raw(raw_tensor, [raw_shape[p] if p >= 0 else 1 for p in out_permutation])
         out_dims = [
-            dims[i]
-            if p >= 0
-            else Dim(
-                kind=dims[i].kind,
-                description="%s_bc_dim1" % (dims[i].description or "unnamed"),
-                dimension=1,
-                auto_generated=True,
+            (
+                dims[i]
+                if p >= 0
+                else Dim(
+                    kind=dims[i].kind,
+                    description="%s_bc_dim1" % (dims[i].description or "unnamed"),
+                    dimension=1,
+                    auto_generated=True,
+                )
             )
             for i, p in enumerate(out_permutation)
         ]
@@ -1300,6 +1302,7 @@ class _TensorMixin(_TensorMixinBase):
         """
         Simpler variant of :func:`copy_compatible_to` which just takes a list of dims,
         and uses simple Dim equality.
+        This adds broadcast dims for any missing dims.
 
         :param dims:
         :return: raw tensor from self with dims permuted and broadcast dims added
@@ -3342,9 +3345,9 @@ class _TensorMixin(_TensorMixinBase):
             max_match_priority = max(dim.match_priority for dim in self.dims)
             return max(
                 matching,
-                key=lambda ax: (max_match_priority + 1)
-                if (self.dims[ax] is other_axis_dim_tag)
-                else self.dims[ax].match_priority,
+                key=lambda ax: (
+                    (max_match_priority + 1) if (self.dims[ax] is other_axis_dim_tag) else self.dims[ax].match_priority
+                ),
             )
 
         other_to_self_mapping = {}

{returnn-1.20240611.113715 → returnn-1.20240614.213333}/returnn/torch/frontend/_backend.py

@@ -1258,7 +1258,7 @@ class TorchBackend(Backend[torch.Tensor]):
         assert all(isinstance(dim, Dim) for dim in axis)
         raw_dims = [source.get_axis_from_description(dim) for dim in axis]
         res_dims = [dim for i, dim in enumerate(source.dims) if i not in raw_dims]
-        correction_factor: Optional[torch.Tensor] = None
+        correction_factor: Union[None, float, Tensor] = None
         if use_mask and any(dim.need_masking() for dim in axis):
             source = source.copy()
             dtype = source.raw_tensor.dtype
@@ -1270,21 +1270,7 @@ class TorchBackend(Backend[torch.Tensor]):
                 mask_value = 0
             elif mode == "mean":
                 mask_value = 0
-                for dim in axis:
-                    if dim.need_masking():
-                        total_num_el = dim.get_dim_value_tensor()
-                        actual_num_el = dim.get_size_tensor()
-                        num_el_reduce_dims = [dim_ for dim_ in axis if dim_ in actual_num_el.dims]
-                        if num_el_reduce_dims:
-                            actual_num_el = rf.reduce_sum(actual_num_el, axis=num_el_reduce_dims)
-                            for dim_ in num_el_reduce_dims:
-                                total_num_el *= dim_.get_dim_value_tensor()
-                        correction_factor_ = rf.cast(total_num_el, source.dtype) / rf.cast(actual_num_el, source.dtype)
-                        correction_factor__ = correction_factor_.copy_compatible_to_dims_raw(res_dims)
-                        if correction_factor is None:
-                            correction_factor = correction_factor__
-                        else:
-                            correction_factor *= correction_factor__
+                correction_factor = rf.masked_fraction_of_shape(axis, inverse=True)
             else:
                 raise NotImplementedError(f"reduce_{mode} not implemented with masking on tensor {source!r}.")
             for dim in axis:
@@ -1302,7 +1288,11 @@ class TorchBackend(Backend[torch.Tensor]):
         else:
             raw_result = func(source.raw_tensor, dim=raw_dims)
         if correction_factor is not None:
-            raw_result *= correction_factor.to(raw_result.device)
+            raw_result *= (
+                correction_factor.copy_compatible_to_dims_raw(res_dims).to(raw_result.device)
+                if isinstance(correction_factor, Tensor)
+                else correction_factor
+            )
         res = Tensor(
             name=f"reduce_{mode}",
             raw_tensor=raw_result,

returnn-1.20240614.213333/returnn/util/math.py

@@ -0,0 +1,28 @@
+"""
+Some mathematical functions, in pure NumPy.
+"""
+
+from __future__ import annotations
+from typing import Union, Dict
+import numpy
+
+
+def next_power_of_two(n: int) -> int:
+    """next power of two, >= n"""
+    return 2 ** (int(n - 1).bit_length())
+
+
+class PiecewiseLinear:
+    """
+    Piecewise linear function.
+    """
+
+    def __init__(self, values: Dict[Union[int, float], Union[int, float]]):
+        self._sorted_items = sorted(values.items())
+        self._sorted_keys = numpy.array([x for x, _ in self._sorted_items])
+        self._sorted_values = numpy.array([y for _, y in self._sorted_items])
+
+    def __call__(self, x: Union[int, float]) -> Union[int, float]:
+        steps = self._sorted_keys
+        values = self._sorted_values
+        return numpy.interp(x, steps, values)

{returnn-1.20240611.113715 → returnn-1.20240614.213333/returnn.egg-info}/PKG-INFO

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

{returnn-1.20240611.113715 → returnn-1.20240614.213333}/tests/test_Util.py

@@ -96,6 +96,26 @@ def test_slice_pad_zeros():
     assert_equal(list(slice_pad_zeros(np.array([1, 2, 3, 4]), begin=2, end=6)), [3, 4, 0, 0])
 
 
+def test_math_PiecewiseLinear():
+    from returnn.util.math import PiecewiseLinear
+
+    eps = 1e-5
+    f = PiecewiseLinear({1: 2, 3: 4, 5: 1})
+    assert_equal(f(0), 2)
+    assert_equal(f(1 - eps), 2)
+    assert_equal(f(1), 2)
+    assert_almost_equal(f(1 + eps), 2, decimal=4)
+    assert_equal(f(2), 3)
+    assert_almost_equal(f(3 - eps), 4, decimal=4)
+    assert_equal(f(3), 4)
+    assert_almost_equal(f(3 + eps), 4, decimal=4)
+    assert_equal(f(4), 2.5)
+    assert_almost_equal(f(5 - eps), 1, decimal=4)
+    assert_equal(f(5), 1)
+    assert_equal(f(5 + eps), 1)
+    assert_equal(f(6), 1)
+
+
 def test_parse_orthography_into_symbols():
     assert_equal(list("hi"), parse_orthography_into_symbols("hi"))
     assert_equal(list(" hello "), parse_orthography_into_symbols(" hello "))

returnn-1.20240611.113715/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240611.113715'
-long_version = '1.20240611.113715+git.8f3d985'

returnn-1.20240611.113715/returnn/util/math.py

@@ -1,11 +0,0 @@
-"""
-Some mathematical functions, in pure NumPy.
-"""
-
-
-from __future__ import annotations
-
-
-def next_power_of_two(n: int) -> int:
-    """next power of two, >= n"""
-    return 2 ** (int(n - 1).bit_length())