returnn 1.20251006.114241__py3-none-any.whl → 1.20251007.115327__py3-none-any.whl

returnn/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20251006.114241
+Version: 1.20251007.115327
 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.20251006.114241'
-long_version = '1.20251006.114241+git.7745ba7'
+version = '1.20251007.115327'
+long_version = '1.20251007.115327+git.70a1d5d'

returnn/datasets/basic.py

@@ -19,6 +19,7 @@ import os
 import math
 import numpy
 import functools
+import types
 from typing import TYPE_CHECKING, Optional, Any, Set, Tuple, Union, Type, Dict, Sequence, List, Callable
 
 from returnn.log import log
@@ -154,7 +155,7 @@ class Dataset:
         self.seq_tags_filter = set(self._load_seq_list_file(seq_list_filter_file)) if seq_list_filter_file else None
         self.unique_seq_tags = unique_seq_tags
         self._seq_order_seq_lens_file = seq_order_seq_lens_file
-        self._seq_order_seq_lens_by_idx = None
+        self._seq_order_seq_lens_by_idx: Optional[Sequence[Union[int, float]]] = None
         # There is probably no use case for combining the two, so avoid potential misconfiguration.
         assert self.partition_epoch == 1 or self.repeat_epoch == 1, (
             "Combining partition_epoch and repeat_epoch is prohibited."
@@ -486,12 +487,8 @@ class Dataset:
         """
         raise NotImplementedError
 
-    def _get_seq_order_seq_lens_by_idx(self, seq_idx):
-        """
-        :param int seq_idx:
-        :rtype: int
-        """
-        if not self._seq_order_seq_lens_by_idx:
+    def _get_seq_order_seq_lens_by_idx(self, seq_idx: int) -> Union[int, float]:
+        if self._seq_order_seq_lens_by_idx is None:
             assert self._seq_order_seq_lens_file
             if self._seq_order_seq_lens_file.endswith(".gz"):
                 import gzip
@@ -502,11 +499,12 @@ class Dataset:
             seq_lens = eval(raw)
             assert isinstance(seq_lens, dict)
             all_tags = self.get_all_tags()
-            self._seq_order_seq_lens_by_idx = [seq_lens[tag] for tag in all_tags]
+            self._seq_order_seq_lens_by_idx = numpy.array([seq_lens[tag] for tag in all_tags])
+        self._get_seq_order_seq_lens_by_idx = self._seq_order_seq_lens_by_idx.__getitem__  # faster
         return self._seq_order_seq_lens_by_idx[seq_idx]
 
     def get_seq_order_for_epoch(
-        self, epoch: Optional[int], num_seqs: int, get_seq_len: Optional[Callable[[int], int]] = None
+        self, epoch: Optional[int], num_seqs: int, get_seq_len: Optional[Callable[[int], Union[int, float]]] = None
     ) -> Sequence[int]:
         """
         Returns the order of the given epoch.
@@ -515,7 +513,7 @@ class Dataset:
 
         :param epoch: for 'random', this determines the random seed
         :param num_seqs:
-        :param get_seq_len: function (originalSeqIdx: int) -> int
+        :param get_seq_len: function (originalSeqIdx: int) -> int|float
         :return: the order for the given epoch. such that seq_idx -> underlying idx
         """
         if epoch is None:
@@ -561,8 +559,9 @@ class Dataset:
             seq_index = range(num_seqs - 1, -1, -1)  # type: Union[range, Sequence[int]]
         elif seq_ordering_method in ["sorted", "sorted_reverse"]:
             assert get_seq_len
-            reverse = -1 if seq_ordering_method == "sorted_reverse" else 1
-            seq_lens = [reverse * get_seq_len(i) for i in range(num_seqs)]
+            seq_lens = _get_seq_len_as_array(get_seq_len, num_seqs)
+            if seq_ordering_method == "sorted_reverse":
+                seq_lens = -seq_lens
             seq_index = numpy.argsort(seq_lens, kind="stable")
         elif seq_ordering_method == "random" or seq_ordering_method.startswith("random:"):
             tmp = seq_ordering_method.split(":", 1)
@@ -628,7 +627,7 @@ class Dataset:
                 nth = 1
             else:
                 nth = int(tmp[1])
-            seq_lens = numpy.array([get_seq_len(i) for i in range(num_seqs)])
+            seq_lens = _get_seq_len_as_array(get_seq_len, num_seqs)
             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)  # type: Union[numpy.ndarray, List[int]]
@@ -1757,3 +1756,19 @@ def set_config_extern_data_from_dataset(config, dataset):
         "extern_data",
         {key: _data_kwargs_from_dataset_key(dataset=dataset, key=key) for key in dataset.get_data_keys()},
     )
+
+
+def _get_seq_len_as_array(get_seq_len: Callable[[int], Union[int, float]], num_seqs: int) -> numpy.ndarray:
+    if num_seqs == 0:
+        return numpy.zeros((0,), dtype=numpy.int32)
+    if isinstance(get_seq_len, (types.BuiltinMethodType, types.MethodWrapperType, types.MethodType)):
+        # Call it once. This might trigger some caching.
+        get_seq_len(0)
+        # Get it again. This might now get us a different (cached) function, e.g. array.__getitem__.
+        get_seq_len = getattr(get_seq_len.__self__, get_seq_len.__name__)
+        assert isinstance(get_seq_len, (types.BuiltinMethodType, types.MethodWrapperType, types.MethodType))
+        obj = get_seq_len.__self__
+        if isinstance(obj, numpy.ndarray) and get_seq_len.__name__ == "__getitem__":
+            assert obj.shape == (num_seqs,)
+            return obj
+    return numpy.array([get_seq_len(i) for i in range(num_seqs)])

returnn/datasets/distrib_files.py

@@ -135,7 +135,7 @@ class DistributeFilesDataset(CachedDataset2):
     def __init__(
         self,
         *,
-        files: Union[List[FileTree], os.PathLike],
+        files: Union[List[FileTree], os.PathLike, Callable[[], List[FileTree]]],
         get_sub_epoch_dataset: Callable[[List[FileTree]], Dict[str, Any]],
         preload_next_n_sub_epochs: int = 1,
         buffer_size: int = 1,
@@ -151,6 +151,7 @@ class DistributeFilesDataset(CachedDataset2):
             can also be specified as a path to a .txt file containing one file per line,
             or a python file containing the repr of a list of arbitrarily nested python objects,
             or a JSON file containing a list of arbitarily nested (JSON) objects.
+            It can also be a callable which returns such a list.
         :param get_sub_epoch_dataset: callable which returns a dataset dict for a given subset of files
         :param preload_next_n_sub_epochs: how many sub epoch datasets to preload
         :param buffer_size: buffer size for each worker, number of seqs to prefetch
@@ -244,6 +245,11 @@ class DistributeFilesDataset(CachedDataset2):
             return
         if isinstance(self.files, list):
             self._files = self.files
+        elif callable(self.files):
+            self._files = self.files()
+            assert isinstance(self._files, list), (
+                f"{self}: callable files {self.files} must return a list, got {type(self._files)}"
+            )
         elif isinstance(self.files, (str, os.PathLike)):
             _, ext = os.path.splitext(self.files)
             assert ext, f"{self}: no file extension on file list file {self.files}"

{returnn-1.20251006.114241.dist-info → returnn-1.20251007.115327.dist-info}/METADATA

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

{returnn-1.20251006.114241.dist-info → returnn-1.20251007.115327.dist-info}/RECORD

@@ -1,9 +1,9 @@
-returnn/PKG-INFO,sha256=khjeqHtAYC68kPEh0ltnHLB14E2_UOObt3MGLvgeMTY,5215
+returnn/PKG-INFO,sha256=nPg5FDephK9Q9lpTYYnNiyWNuWX4iqkbzK4OJytg_64,5215
 returnn/__init__.py,sha256=biBtRsM0WZ406vShaeH-9WFoqJ8XwTbn6g0EeFJ7l8E,1012
 returnn/__main__.py,sha256=lHyZcu_0yc9f7Vf_Kfdy9PmeU0T76XVXnpalHi5WKro,31740
 returnn/__old_mod_loader__.py,sha256=nvsNY-xELdS_IPNkv66Q9Rmvg4dbGW0-EBRDcCmctos,7654
 returnn/__setup__.py,sha256=22kQn2fh11iPM0hLb2Fy5sLmoU1JGvmDxXRYuRgQkwU,4659
-returnn/_setup_info_generated.py,sha256=MFwUI_cce-3opoV_eKfpRd891cRwBo6RMbTGriz-HTg,77
+returnn/_setup_info_generated.py,sha256=3Oo-LSdu1OoWfV1GeXf8IFcg75AqFRPpH9erSCeRGPA,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
@@ -13,11 +13,11 @@ returnn/native_op.py,sha256=4_NnvfNxsM8GE_FsD6yOg6PZegqIdtJ3Sl1GdBWmFvg,244424
 returnn/pretrain.py,sha256=MHiXJZqkQFmDVyaYsGpd_Acv20wxl7Pr6s6qJzAT2FI,22648
 returnn/datasets/__init__.py,sha256=PvDlfDOaaopIeUIt0OSvHD2eHZkdkyE-sjMXf35EH5U,390
 returnn/datasets/audio.py,sha256=Gmj7a08dnvYh7Z-G1TNapz42L50AIcDE9JeIZaO1s1M,23334
-returnn/datasets/basic.py,sha256=_42fQztTZq7jNQrWdFBwulB1bNta17LOTyrD8XJ-7_E,73089
+returnn/datasets/basic.py,sha256=e-hrhJjLq01JT3dYxeeRhw4z1DQC5lxglpw4RscD57U,74136
 returnn/datasets/bundle_file.py,sha256=KQNrS1MSf-4_idlK0c0KFwON-f5sEK0sWU15WpoMYpE,2380
 returnn/datasets/cached.py,sha256=RyefRjSDdp-HveK-2vLy2C6BIHcpqQ_lNvUKlIa4QAI,25412
 returnn/datasets/cached2.py,sha256=oJOq2lWRQpxm6kyUKW1w5qZBd4kdKEpwM7KY_QnXbq4,11922
-returnn/datasets/distrib_files.py,sha256=-WNVhtvdJFP3L9Meh33oTSYc0FJSvF40mJ5UI_vJbSE,30233
+returnn/datasets/distrib_files.py,sha256=srTieLP02kCepAwZ6Y9p20cqB8nAlVJWbSAoOPna9ik,30567
 returnn/datasets/generating.py,sha256=Qb7V94N_GfL2pZPxWS5PmzszoVXXKzuUmsHuW3dmVbc,99556
 returnn/datasets/hdf.py,sha256=v5sjBenURR9Z-g7AQ9tsL84yDSye5RtbLpym3M6HSDE,67833
 returnn/datasets/lm.py,sha256=rQ3jV43lSnlGkKu7m5jTTH7aK0BOMXQocsHfJ8OGec8,99950
@@ -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.20251006.114241.dist-info/LICENSE,sha256=ywBD_U2aD4vpuoIgNAsjIGBYydl0tVKll3De0Z8s77c,11041
-returnn-1.20251006.114241.dist-info/METADATA,sha256=khjeqHtAYC68kPEh0ltnHLB14E2_UOObt3MGLvgeMTY,5215
-returnn-1.20251006.114241.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
-returnn-1.20251006.114241.dist-info/top_level.txt,sha256=Lsn4WZc5Pbfk0-xDQOgnFCxOoqxL4CyeM3N1TFbJncw,8
-returnn-1.20251006.114241.dist-info/RECORD,,
+returnn-1.20251007.115327.dist-info/LICENSE,sha256=ywBD_U2aD4vpuoIgNAsjIGBYydl0tVKll3De0Z8s77c,11041
+returnn-1.20251007.115327.dist-info/METADATA,sha256=nPg5FDephK9Q9lpTYYnNiyWNuWX4iqkbzK4OJytg_64,5215
+returnn-1.20251007.115327.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
+returnn-1.20251007.115327.dist-info/top_level.txt,sha256=Lsn4WZc5Pbfk0-xDQOgnFCxOoqxL4CyeM3N1TFbJncw,8
+returnn-1.20251007.115327.dist-info/RECORD,,