returnn 1.20241113.165719__tar.gz → 1.20241115.151509__tar.gz

{returnn-1.20241113.165719 → returnn-1.20241115.151509}/PKG-INFO

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

returnn-1.20241115.151509/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20241115.151509'
+long_version = '1.20241115.151509+git.57080cf'

{returnn-1.20241113.165719 → returnn-1.20241115.151509}/returnn/datasets/basic.py

@@ -120,7 +120,7 @@ class Dataset:
             not all datasets support this option.
         :param None|int|dict|NumbersDict|(dict,dict) context_window: will add this context for each chunk
         :param None|str|int|(int,int)|dict|(dict,dict)|function chunking: "chunk_size:chunk_step"
-        :param str seq_ordering: "batching"-option in config. e.g. "default", "sorted" or "random".
+        :param str|function seq_ordering: "batching"-option in config. e.g. "default", "sorted" or "random".
             See self.get_seq_order_for_epoch() for more details.
         :param int|None fixed_random_seed: for the shuffling, e.g. for seq_ordering='random'.
             otherwise epoch will be used.
@@ -511,34 +511,50 @@ class Dataset:
         if self._seq_order_seq_lens_file:
             get_seq_len = self._get_seq_order_seq_lens_by_idx
 
-        if self.seq_ordering == "default":
+        seq_ordering_method = self.seq_ordering
+        if callable(seq_ordering_method):
+            seq_ordering_method_ = seq_ordering_method(epoch=epoch, **get_fwd_compat_kwargs())
+            if not isinstance(seq_ordering_method_, str):
+                raise TypeError(
+                    f"seq_ordering function {seq_ordering_method} should return a str,"
+                    f" got {seq_ordering_method_!r} (type {type(seq_ordering_method_).__name__})"
+                )
+            seq_ordering_method = seq_ordering_method_
+        elif isinstance(seq_ordering_method, str):
+            pass  # fine, checked below
+        else:
+            raise TypeError(
+                f"seq_ordering {seq_ordering_method!r} unexpected type {type(seq_ordering_method).__name__}"
+            )
+
+        if seq_ordering_method == "default":
             seq_index = range(num_seqs)
-        elif self.seq_ordering.startswith("default_every_n:"):
+        elif seq_ordering_method.startswith("default_every_n:"):
             # This order is useful if you have "initial_state": "keep_over_epoch",
             # where num == max_seqs, batch_size = inf, max_seq_len = inf, chunking = None.
-            _, num = self.seq_ordering.split(":")
+            _, num = seq_ordering_method.split(":")
             num = int(num)
             seq_index = numpy.arange(num_seqs // num, dtype="int64").repeat(num)
             for i in range(1, num):
                 seq_index[i::num] += i * (num_seqs // num)
-        elif self.seq_ordering == "reverse":
+        elif seq_ordering_method == "reverse":
             seq_index = range(num_seqs - 1, -1, -1)  # type: Union[range, typing.Sequence[int]]
-        elif self.seq_ordering in ["sorted", "sorted_reverse"]:
+        elif seq_ordering_method in ["sorted", "sorted_reverse"]:
             assert get_seq_len
-            reverse = -1 if self.seq_ordering == "sorted_reverse" else 1
+            reverse = -1 if seq_ordering_method == "sorted_reverse" else 1
             seq_lens = [reverse * get_seq_len(i) for i in range(num_seqs)]
             seq_index = numpy.argsort(seq_lens, kind="stable")
-        elif self.seq_ordering.startswith("random"):
-            tmp = self.seq_ordering.split(":")
+        elif seq_ordering_method.startswith("random"):
+            tmp = seq_ordering_method.split(":")
             nth = int(tmp[1]) if len(tmp) > 1 else 1
             # Keep this deterministic! Use fixed seed.
             rnd_seed = self._get_random_seed_for_epoch(epoch=epoch, num_epochs_fixed=nth)
             random_generator = numpy.random.RandomState(rnd_seed)
             seq_index = random_generator.permutation(num_seqs)
-        elif self.seq_ordering.startswith("sort_bin_shuffle"):
+        elif seq_ordering_method.startswith("sort_bin_shuffle"):
             # Shuffle seqs, sort by length, and shuffle bins (then shuffle seqs within each bin if sort_bin_shuffle_x2).
             assert get_seq_len
-            tmp = self.seq_ordering.split(":")[1:]
+            tmp = seq_ordering_method.split(":")[1:]
             # Keep this deterministic! Use fixed seed.
             if len(tmp) <= 1:
                 nth = 1
@@ -562,13 +578,13 @@ class Dataset:
                     part = seq_index[i * len(seq_index) // bins :][:]
                 else:
                     part = seq_index[i * len(seq_index) // bins : (i + 1) * len(seq_index) // bins][:]
-                if self.seq_ordering.startswith("sort_bin_shuffle_x2"):
+                if seq_ordering_method.startswith("sort_bin_shuffle_x2"):
                     random_generator.shuffle(part)  # Shuffle within the bin.
                 out_index.append(part)
             seq_index = numpy.concatenate(out_index)
-        elif self.seq_ordering.startswith("laplace"):
+        elif seq_ordering_method.startswith("laplace"):
             assert get_seq_len
-            tmp = self.seq_ordering.split(":")[1:]
+            tmp = seq_ordering_method.split(":")[1:]
             if len(tmp) == 0:
                 bins = 2
             else:
@@ -595,7 +611,7 @@ class Dataset:
                 out_index.append(part)
             seq_index = numpy.concatenate(out_index)
         else:
-            assert False, "invalid batching specified: " + self.seq_ordering
+            raise ValueError(f"invalid seq_ordering {seq_ordering_method!r}")
 
         if self.unique_seq_tags:
             # Note: This is as generic as possible, but requires that get_all_tags is implemented.

{returnn-1.20241113.165719 → returnn-1.20241115.151509}/returnn/datasets/multi_proc.py

@@ -66,6 +66,7 @@ class MultiProcDataset(CachedDataset2):
             )
         self._sharding_method = sharding_method
         self._data_keys = None
+        self._data_dtypes: Optional[Dict[str, str]] = None
         self._num_seqs = None
         self._total_num_seqs = None
         self._all_tags = None
@@ -83,6 +84,8 @@ class MultiProcDataset(CachedDataset2):
             self.num_outputs = _meta_info_cache["num_outputs"]
             self._total_num_seqs = _meta_info_cache["total_num_seqs"]
             self.labels = _meta_info_cache["labels"]
+            self._data_keys = _meta_info_cache["data_keys"]
+            self._data_dtypes = _meta_info_cache["data_dtypes"]
 
     def initialize(self):
         """init"""
@@ -100,6 +103,8 @@ class MultiProcDataset(CachedDataset2):
             "labels": self.labels,
             "total_num_seqs": self._total_num_seqs,
             "all_tags": self._all_tags,
+            "data_keys": self._data_keys,
+            "data_dtypes": self._data_dtypes,
         }
 
     def _lazy_init(self):
@@ -177,6 +182,26 @@ class MultiProcDataset(CachedDataset2):
         msg, self.labels = self._seq_order_proc_parent_conn.recv()
         assert msg == "labels"
 
+    def _lazy_init_data_keys(self):
+        if self._data_keys is not None:
+            return
+
+        self._lazy_init()  # indempotent, will not run twice if not needed
+
+        self._seq_order_proc_parent_conn.send(("get_data_keys", {}))
+        msg, self._data_keys = self._seq_order_proc_parent_conn.recv()
+        assert msg == "data_keys"
+
+    def _lazy_init_data_dtypes(self):
+        if self._data_dtypes is not None:
+            return
+
+        self._lazy_init()  # indempotent, will not run twice if not needed
+
+        self._seq_order_proc_parent_conn.send(("get_data_dtypes", {}))
+        msg, self._data_dtypes = self._seq_order_proc_parent_conn.recv()
+        assert msg == "data_dtypes"
+
     def __del__(self):
         if self._worker_procs:
             got_exception = False
@@ -298,6 +323,14 @@ class MultiProcDataset(CachedDataset2):
                     except NotImplementedError as exc:
                         all_tags = NotImplementedError(f"{exc} in {dataset}")
                     parent_conn.send(("all_tags", all_tags))
+                elif msg == "get_data_keys":
+                    assert dataset
+                    data_keys = dataset.get_data_keys()
+                    parent_conn.send(("data_keys", data_keys))
+                elif msg == "get_data_dtypes":
+                    assert dataset
+                    data_keys = dataset.get_data_keys()
+                    parent_conn.send(("data_dtypes", {k: dataset.get_data_dtype(k) for k in data_keys}))
                 elif msg == "init_seq_order":
                     if dataset is None:
                         dataset = init_dataset(dataset_dict)
@@ -444,3 +477,24 @@ class MultiProcDataset(CachedDataset2):
         super().finish_epoch(free_resources=free_resources)
         for worker_parent_conn in self._worker_parent_conns:
             worker_parent_conn.send(("finish_epoch", {"free_resources": free_resources}))
+
+    def get_data_keys(self) -> List[str]:
+        """data keys"""
+        if self._data_keys is None:
+            self._lazy_init_data_keys()
+            assert self._data_keys is not None
+        return self._data_keys
+
+    def get_data_dtype(self, key: str) -> str:
+        """:return: dtype of `key`"""
+        if self._data_dtypes is None:
+            self._lazy_init_data_dtypes()
+            assert self._data_dtypes is not None
+        return self._data_dtypes[key]
+
+    def is_data_sparse(self, key: str) -> bool:
+        """:return: whether `key` is sparse"""
+        if self.num_outputs is None:
+            self._lazy_init()
+            assert self.num_outputs is not None
+        return super().is_data_sparse(key)

{returnn-1.20241113.165719 → returnn-1.20241115.151509}/returnn/util/math.py

@@ -3,9 +3,10 @@ Some mathematical functions, in pure NumPy.
 """
 
 from __future__ import annotations
-from typing import Union, Optional, Sequence, Dict
+from typing import Union, Optional, Any, Sequence, Dict
 import numpy
 import hashlib
+import bisect
 
 
 def ceil_div(a: int, b: int) -> int:
@@ -80,6 +81,74 @@ class PiecewiseLinear:
         return numpy.interp(x, steps, values)
 
 
+class StepFunction:
+    """
+    Step function. (Piecewise constant function.)
+
+    It will find the nearest x value and return the corresponding y value.
+    The corresponding y value can be anything here, not just numbers.
+    """
+
+    def __init__(
+        self,
+        values: Dict[Union[int, float], Any],
+        *,
+        kw_name: Optional[str] = None,
+        ignore_other_kwargs: bool = False,
+    ):
+        """
+        :param values: dict x -> y. Everything between the x values is the value of the nearest x.
+        :param kw_name: keyword argument name to use in the __call__. Other keyword arguments are ignored.
+        :param ignore_other_kwargs: if True, ignore other keyword arguments in the __call__.
+        """
+        assert values, f"{self}: No values provided"
+        self._values = values
+        self._sorted_keys = sorted(values.keys())
+        self._kw_name = kw_name
+        self._ignore_other_kwargs = ignore_other_kwargs
+
+    def __getstate__(self):
+        # Note: I was implementing __getnewargs_ex__, but we cannot use this because of this Sisyphus bug:
+        # https://github.com/rwth-i6/sisyphus/issues/215
+        kwargs = {"values": self._values}
+        if self._kw_name is not None:
+            kwargs["kw_name"] = self._kw_name
+        if self._ignore_other_kwargs:
+            kwargs["ignore_other_kwargs"] = True
+        return kwargs
+
+    def __setstate__(self, state):
+        self.__init__(**state)
+
+    def __repr__(self) -> str:
+        kwargs = self.__getstate__()
+        values = kwargs.pop("values")
+        all_args_str = ", ".join([repr(values)] + [f"{k}={v!r}" for k, v in kwargs.items()])
+        return f"{self.__class__.__name__}({all_args_str})"
+
+    def __call__(self, *args: Union[int, float], **kwargs) -> Union[int, float]:
+        if self._kw_name:
+            if args:
+                raise TypeError(f"{self}: Expected zero positional arguments, got {args!r}")
+            x = kwargs.pop(self._kw_name, None)
+        else:
+            if len(args) != 1:
+                raise TypeError(f"{self}: Expected one positional argument, got {args!r}")
+            x = args[0]
+        if not self._ignore_other_kwargs:
+            if kwargs:
+                raise TypeError(f"{self}: Unexpected keyword arguments: {kwargs!r}")
+        assert x is not None
+        keys = self._sorted_keys
+        assert keys, f"{self}: No keys provided"
+        i = bisect.bisect_left(keys, x)
+        # keys[:i] have e < x, and all e in keys[i:] have e >= x
+        # keys[i] is the smallest e >= x
+        if i >= len(keys) or (i > 0 and (x - keys[i - 1]) <= (keys[i] - x)):
+            i -= 1
+        return self._values[keys[i]]
+
+
 def simplify_and_format_number(n: Union[int, float]) -> str:
     """Format a number, removing trailing zeros and the decimal point if it is an integer"""
     if isinstance(n, (int, float)):

{returnn-1.20241113.165719 → returnn-1.20241115.151509}/returnn.egg-info/PKG-INFO

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

{returnn-1.20241113.165719 → returnn-1.20241115.151509}/tests/test_Util.py

@@ -146,6 +146,18 @@ def test_math_PiecewiseLinear_kwargs():
     assert f(epoch_continuous=0, seq_idx=123) == 2
 
 
+def test_math_StepFunction_kwargs():
+    from returnn.util.math import StepFunction
+
+    f = StepFunction({1: "yes", 3: "no"}, kw_name="epoch", ignore_other_kwargs=True)
+    assert f(epoch=0) == "yes"
+    assert f(epoch=1) == "yes"
+    assert f(epoch=2) == "yes"  # boundary, assume left wins
+    assert f(epoch=2.1) == "no"
+    assert f(epoch=3) == "no"
+    assert f(epoch=100, seq_idx=123) == "no"
+
+
 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.20241113.165719/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20241113.165719'
-long_version = '1.20241113.165719+git.138d3e9'