returnn 1.20240610.120719__tar.gz → 1.20240613.164354__tar.gz

{returnn-1.20240610.120719/returnn.egg-info → returnn-1.20240613.164354}/PKG-INFO

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

returnn-1.20240613.164354/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240613.164354'
+long_version = '1.20240613.164354+git.f006453'

{returnn-1.20240610.120719 → returnn-1.20240613.164354}/returnn/datasets/basic.py

@@ -1367,7 +1367,7 @@ def get_dataset_class(name: Union[str, Type[Dataset]]) -> Optional[Type[Dataset]
         "raw_wav",
         "map",
         "multi_proc",
-        "concat_files",
+        "distrib_files",
     ]
     for mod_name in mod_names:
         mod = import_module("returnn.datasets.%s" % mod_name)

returnn-1.20240610.120719/returnn/datasets/concat_files.py → returnn-1.20240613.164354/returnn/datasets/distrib_files.py

@@ -1,5 +1,5 @@
 """
-:class:`ConcatFilesDataset `
+:class:`DistributeFilesDataset`
 
 https://github.com/rwth-i6/returnn/issues/1519
 """
@@ -23,19 +23,21 @@ from multiprocessing.connection import Connection as mpConnection
 _mp = NonDaemonicSpawnContext(process_pre_init_func=SubProcCopyGlobalConfigPreInitFunc())
 
 
-__all__ = ["ConcatFilesDataset"]
+__all__ = ["DistributeFilesDataset"]
 
 Filename = str
 FileTree = Union[Filename, Tuple["FileTree", ...], Dict[Any, "FileTree"], List["FileTree"]]
 
 
-class ConcatFilesDataset(CachedDataset2):
+class DistributeFilesDataset(CachedDataset2):
     """
-    This is similar to :class:`ConcatDataset`, but instead of concatenating datasets,
-    we distribute files over subepochs,
-    and then create a sub dataset for every sub epoch
-    for a given subset of the files
-    via the given ``get_sub_epoch_dataset``.
+    Dataset that distributes files over subepochs and then creates a
+    sub dataset for every sub epoch for a given (random) subset of the files.
+    The sub dataset is user-defined via a function ``get_sub_epoch_dataset``.
+    Thus, this dataset wraps the sub datasets.
+
+    It is conceptually very similar to :class:`ConcatDataset` in the sense
+    that it concatenates all the sub datasets together to form one larger dataset.
 
     This scheme allows to shuffle over the files,
     which makes shuffling much more efficient over a large dataset
@@ -70,7 +72,7 @@ class ConcatFilesDataset(CachedDataset2):
           }
 
         train = {
-          "class": "ConcatFilesDataset",
+          "class": "DistributeFilesDataset",
           "files": [
             "/nfs/big_data_1.hdf",
             ...
@@ -105,7 +107,7 @@ class ConcatFilesDataset(CachedDataset2):
           }
 
         train = {
-          "class": "ConcatFilesDataset",
+          "class": "DistributeFilesDataset",
           "files": [
             ("/nfs/alignment_1.hdf", "/nfs/features_1.hdf"),
             ...
@@ -321,7 +323,7 @@ class ConcatFilesDataset(CachedDataset2):
         # [[1,1], [78], [120], []] or [[1,1,78], [120], [], []].
         # Or consider [5,5]+[10]*7, partition_epoch=5, which has avg size 16.
         # A simple algorithm could end up with [[5,5,10], [10,10], [10,10], [10,10], []].
-        # See test_ConcatFilesDataset_get_files_per_sub_epochs for some test cases.
+        # See test_DistributeFilesDataset_get_files_per_sub_epochs for some test cases.
         assert len(files_order) >= partition_epoch
         files_per_sub_epochs = [[] for _ in range(partition_epoch)]
         assert len(files_per_sub_epochs) == partition_epoch

{returnn-1.20240610.120719 → returnn-1.20240613.164354}/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.20240610.120719 → returnn-1.20240613.164354}/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.20240610.120719 → returnn-1.20240613.164354}/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.20240610.120719 → returnn-1.20240613.164354}/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.20240610.120719 → returnn-1.20240613.164354/returnn.egg-info}/PKG-INFO

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

{returnn-1.20240610.120719 → returnn-1.20240613.164354}/returnn.egg-info/SOURCES.txt

@@ -99,7 +99,7 @@ returnn/datasets/basic.py
 returnn/datasets/bundle_file.py
 returnn/datasets/cached.py
 returnn/datasets/cached2.py
-returnn/datasets/concat_files.py
+returnn/datasets/distrib_files.py
 returnn/datasets/generating.py
 returnn/datasets/hdf.py
 returnn/datasets/lm.py

{returnn-1.20240610.120719 → returnn-1.20240613.164354}/tests/test_Dataset.py

@@ -623,13 +623,13 @@ def test_MapDatasetWrapper():
     assert res.features["data"].shape == (5, 3)
 
 
-def test_ConcatFilesDataset_get_files_per_sub_epochs():
-    from returnn.datasets.concat_files import ConcatFilesDataset
+def test_DistributeFilesDataset_get_files_per_sub_epochs():
+    from returnn.datasets.distrib_files import DistributeFilesDataset
 
     def _test(sizes: List[int], partition_epoch: int, expected: List[List[int]]):
         files = [f"file-{i}" for i in range(len(sizes))]
         file_sizes = {f: s for f, s in zip(files, sizes)}
-        res = ConcatFilesDataset._get_files_per_sub_epochs(
+        res = DistributeFilesDataset._get_files_per_sub_epochs(
             partition_epoch=partition_epoch, file_sizes=file_sizes, files_order=files
         )
         assert all(res) and len(res) == partition_epoch
@@ -643,8 +643,8 @@ def test_ConcatFilesDataset_get_files_per_sub_epochs():
     _test([5, 5] + [10] * 7, 5, [[5, 5, 10], [10, 10], [10, 10], [10], [10]])
 
 
-def test_ConcatFilesDataset():
-    from returnn.datasets.concat_files import ConcatFilesDataset
+def test_DistributeFilesDataset():
+    from returnn.datasets.distrib_files import DistributeFilesDataset
     from test_HDFDataset import generate_hdf_from_other
 
     # Create a few HDF files such that we can easily verify the data later.
@@ -686,7 +686,7 @@ def test_ConcatFilesDataset():
         total_num_seqs += 1
     assert total_num_seqs == num_hdf_files * num_seqs
 
-    # Test to load via ConcatFilesDataset.
+    # Test to load via DistributeFilesDataset.
 
     def _get_sub_epoch_dataset(files_subepoch: List[str]) -> Dict[str, Any]:
         return {"class": "HDFDataset", "files": files_subepoch, "seq_ordering": "default"}
@@ -695,13 +695,13 @@ def test_ConcatFilesDataset():
     assert num_hdf_files % partition_epoch == 0  # just for easier testing here
     concat_dataset = init_dataset(
         {
-            "class": "ConcatFilesDataset",
+            "class": "DistributeFilesDataset",
             "files": hdf_files,
             "get_sub_epoch_dataset": _get_sub_epoch_dataset,
             "partition_epoch": partition_epoch,
         }
     )
-    assert isinstance(concat_dataset, ConcatFilesDataset)
+    assert isinstance(concat_dataset, DistributeFilesDataset)
     assert concat_dataset.get_data_keys() == ["classes"]
     num_hdfs_per_part = num_hdf_files // partition_epoch
     global_seq_idx = 0

returnn-1.20240610.120719/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240610.120719'
-long_version = '1.20240610.120719+git.cc96aee'