returnn 1.20250812.202710__py3-none-any.whl → 1.20250817.33823__py3-none-any.whl

returnn/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20250812.202710
+Version: 1.20250817.33823
 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.20250812.202710'
-long_version = '1.20250812.202710+git.6c611de'
+version = '1.20250817.033823'
+long_version = '1.20250817.033823+git.4f41ff3'

returnn/datasets/distrib_files.py

@@ -6,6 +6,7 @@ https://github.com/rwth-i6/returnn/issues/1519
 
 from __future__ import annotations
 
+from collections import deque
 from typing import Union, Optional, Any, Callable, Sequence, Tuple, List, Dict
 import os
 import sys
@@ -636,7 +637,7 @@ def _worker_proc_loop(
     dataset = init_dataset(dataset_dict)
 
     got_init_seq_order = False
-    cache: List[DatasetSeq] = []
+    cache: deque[DatasetSeq] = deque()
     next_seq_idx = 0
 
     # noinspection PyShadowingNames
@@ -701,7 +702,7 @@ def _worker_proc_loop(
             elif msg == "get_data_seq":
                 seq_idx = kwargs["seq_idx"]
                 while cache and cache[0].seq_idx < seq_idx:
-                    cache.pop(0)
+                    cache.popleft()
                 res = _get(seq_idx)
                 parent_conn.send(("data_seq", res))
             elif msg == "init_seq_order":
@@ -714,7 +715,7 @@ def _worker_proc_loop(
                 parent_conn.send(("num_seqs", num_seqs))
                 got_init_seq_order = True
                 next_seq_idx = 0
-                cache[:] = []
+                cache.clear()
             else:
                 raise Exception(f"unknown msg {msg!r}")
     except KeyboardInterrupt:  # when parent dies

returnn/datasets/multi_proc.py

@@ -3,6 +3,7 @@ Multi-processing dataset
 """
 
 from __future__ import annotations
+from collections import deque
 from typing import Optional, Any, Dict, List
 import sys
 import gc
@@ -234,7 +235,7 @@ class MultiProcDataset(CachedDataset2):
         dataset: Optional[Dataset] = None
 
         got_init_seq_order = False
-        cache = []  # type: List[DatasetSeq]
+        cache: deque[DatasetSeq] = deque()
         next_seq_idx = 0
 
         # noinspection PyShadowingNames
@@ -299,7 +300,7 @@ class MultiProcDataset(CachedDataset2):
                 elif msg == "get_data_seq":
                     seq_idx = kwargs["seq_idx"]
                     while cache and cache[0].seq_idx < seq_idx:
-                        cache.pop(0)
+                        cache.popleft()
                     res = _get(seq_idx)
                     parent_conn.send(("data_seq", res))
                 elif msg == "init":
@@ -372,11 +373,11 @@ class MultiProcDataset(CachedDataset2):
                         raise ValueError(f"{MultiProcDataset.__name__}: unknown sharding_method: {sharding_method}")
                     got_init_seq_order = True
                     next_seq_idx = 0
-                    cache[:] = []
+                    cache.clear()
                 elif msg == "finish_epoch":
                     got_init_seq_order = False
                     next_seq_idx = 0
-                    cache[:] = []
+                    cache.clear()
                     if dataset:
                         dataset.finish_epoch(**kwargs)
                     if kwargs["free_resources"]:

returnn/frontend/_backend.py

@@ -1153,6 +1153,13 @@ class Backend(Generic[T]):
         """
         raise NotImplementedError
 
+    @staticmethod
+    def random_choice_with_replacement(dims: Sequence[Dim], *, probs: Tensor, axis: Dim) -> Tensor:
+        """
+        random choice with replacement. See `rf.random_choice_with_replacement` for details.
+        """
+        raise NotImplementedError
+
     @staticmethod
     def masked_select(
         tensor: Tensor, *, mask: Tensor, dims: Sequence[Dim], out_dim: Optional[Dim] = None

returnn/frontend/rand.py

@@ -65,6 +65,7 @@ __all__ = [
     "random_normal",
     "random_truncated_normal",
     "random_choice_without_replacement",
+    "random_choice_with_replacement",
 ]
 
 
@@ -379,3 +380,16 @@ def random_choice_without_replacement(
     scores = log_probs + scores_random_sample
     _, indices, _ = rf.top_k(scores, k_dim=num_samples_dim, axis=axis)
     return indices
+
+
+def random_choice_with_replacement(dims: Sequence[Dim], *, probs: Tensor, axis: Dim) -> Tensor:
+    """
+    Randomly sample with replacement.
+
+    :param dims: {common_dims..., new_dims...}. Defines how many samples to draw. Defines the output shape.
+    :param probs: {common_dims..., axis}
+    :param axis: feature axis, where to sample from
+    :return: random indices shape dims -> axis.
+    """
+    # noinspection PyProtectedMember
+    return probs._raw_backend.random_choice_with_replacement(dims=dims, probs=probs, axis=axis)

returnn/torch/frontend/_backend.py

@@ -1718,6 +1718,28 @@ class TorchBackend(Backend[torch.Tensor]):
             )
         return out
 
+    @staticmethod
+    def random_choice_with_replacement(dims: Sequence[Dim], *, probs: Tensor, axis: Dim) -> Tensor:
+        """random choice with replacement"""
+        assert all(d == axis or d in dims for d in probs.dims), (
+            f"random_choice_with_replacement: dims {dims} not compatible with probs {probs} and axis {axis}"
+        )
+        common_dims = [d for d in dims if d in probs.dims]
+        assert axis not in common_dims
+        probs = probs.copy_transpose(common_dims + [axis])
+        non_common_dims = [d for d in dims if d not in common_dims]
+        num_samples = prod([d.get_dim_value() for d in non_common_dims])
+        if len(common_dims) >= 2:
+            probs, flat_common_dim = rf.merge_dims(probs, dims=common_dims)
+        out_raw = torch.multinomial(probs.raw_tensor, num_samples=num_samples, replacement=True)
+        out_raw = out_raw.reshape(
+            [d.get_dim_value() for d in common_dims] + [d.get_dim_value() for d in non_common_dims]
+        )
+        out = rf.convert_to_tensor(out_raw, dims=common_dims + non_common_dims, sparse_dim=axis)
+        out = out.copy_transpose(dims)
+        out.name = "random_choice_with_replacement"
+        return out
+
     @staticmethod
     def masked_select(
         tensor: Tensor, *, mask: Tensor, dims: Sequence[Dim], out_dim: Optional[Dim] = None

{returnn-1.20250812.202710.dist-info → returnn-1.20250817.33823.dist-info}/METADATA

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

{returnn-1.20250812.202710.dist-info → returnn-1.20250817.33823.dist-info}/RECORD

@@ -1,9 +1,9 @@
-returnn/PKG-INFO,sha256=rpwjyEcl0A4JGcT_ukgciH0uSWUHPI8REBpNnhAaPjs,5215
+returnn/PKG-INFO,sha256=B57ObmyaxFzp_TjsXhRqtVvzxZWq3hS54vw34tZHZsQ,5214
 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=GR2tWQ1k_EaJwxITWrntrJVKLZr3eU5Omp7yXewSGGQ,77
+returnn/_setup_info_generated.py,sha256=BSqlencj66qyQkfuBO7fze_8cqW9p9M0Si-qLIGy1-k,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
@@ -17,13 +17,13 @@ returnn/datasets/basic.py,sha256=_42fQztTZq7jNQrWdFBwulB1bNta17LOTyrD8XJ-7_E,730
 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=TDkuZL_1c-QJyCzGw56alLrWmLRnppNFLtPhqzXm7Tc,30194
+returnn/datasets/distrib_files.py,sha256=MA4BuqJEyen-3aA_naS9RWOHrP3OjVjP63VeWfLWigs,30233
 returnn/datasets/generating.py,sha256=9U_w6URIrv-Rb-hDbPOzYW9qYXzJbw32N6G268IKyoM,99833
 returnn/datasets/hdf.py,sha256=v5sjBenURR9Z-g7AQ9tsL84yDSye5RtbLpym3M6HSDE,67833
 returnn/datasets/lm.py,sha256=rQ3jV43lSnlGkKu7m5jTTH7aK0BOMXQocsHfJ8OGec8,99950
 returnn/datasets/map.py,sha256=kOBJVZmwDhLsOplzDNByIfa0NRSUaMo2Lsy36lBvxrM,10907
 returnn/datasets/meta.py,sha256=6XPPxhiNSxWw9Hu5Z6wG8dD9Zk82FqiI-k9HGQSTKgw,95658
-returnn/datasets/multi_proc.py,sha256=aVjsLt2qjHnHOrEYCgIPCwNYE-f1fiGP6eZ8NGAr3A4,22583
+returnn/datasets/multi_proc.py,sha256=3kZNXbwfuEdQrlOaiY-XnADrQrz7QOqZIQTZUwIWuS8,22614
 returnn/datasets/normalization_data.py,sha256=J3njQCMvWAbIAVPepO2L_Xdau9eWYB7Zyd6STeGzTbc,14615
 returnn/datasets/numpy_dump.py,sha256=wl8bKIKAlff2HPJPtuu5wBg3TLOf16d2wLVB4lLAwTM,5158
 returnn/datasets/postprocessing.py,sha256=6SfT58BxbHYO2QlGzOgIV04Zqkp-kl0B85168DQaB9A,24060
@@ -75,7 +75,7 @@ returnn/extern/graph_editor/subgraph.py,sha256=q9o0zVBLDrTIidaXg5WG5daDW0mLbwv2J
 returnn/extern/graph_editor/transform.py,sha256=qMGSenpbAnGqdG6QP6iWjlm6_ccySYJaZKOoAj1dbOM,29348
 returnn/extern/graph_editor/util.py,sha256=HfRbyQPmQ6_n5-O-096n0KeJtllQXFtaurpeJS_URZ0,18706
 returnn/frontend/__init__.py,sha256=2aS7nbxXniIrBp2DODl0xN0f3IJ_dX4Bi9ZlR7W5_DE,1472
-returnn/frontend/_backend.py,sha256=pAnVAbZhIGKD-10tp0Mx7AO1GZNghYu7AVAPhiimN-k,50471
+returnn/frontend/_backend.py,sha256=39l5MC1DaT0MPklMM8HXAW9nqisIIZQ9g2QSHOOtPQE,50741
 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
@@ -110,7 +110,7 @@ returnn/frontend/parameter.py,sha256=zvrkhSYC1c_O9kVwgHvOtOnWNurl5J28lkS0i1LQpWU
 returnn/frontend/parametrizations.py,sha256=ptNgBw5IiPXVpB3QGse7AGAhdXp8X1rCqYUl2Mae8aI,2876
 returnn/frontend/parametrize.py,sha256=VhgTEP7ehON950Q4bkCy8rvg9641moEKAXn0XzomK6E,7216
 returnn/frontend/piecewise_linear.py,sha256=TdL6wzop8P1dcIZwkEbJFvSUZSI1cbhS3XKzlWQkEVI,1964
-returnn/frontend/rand.py,sha256=Levgf5VtOOBKDSgz0869Jf3VW4BWxYZuRXsa_fOxNI4,12969
+returnn/frontend/rand.py,sha256=2x7AHSYH_tZkzTk_q3t3GA_yYRNeKsVbJjw2InqSGDk,13542
 returnn/frontend/rec.py,sha256=6YSsSG7fdtfvvg24vmexSg8R2aVCcKHBdGLh-Mgn9Co,8037
 returnn/frontend/reduce.py,sha256=gRSvBJZNHa757IqBxGw4hu5eiO3pjie_ptEwUXHLSCs,10340
 returnn/frontend/run_ctx.py,sha256=yyOMUCKTOe19C4z2Nfly4YCLBmQ9ihip6nGrkW-Y6qg,23789
@@ -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=a9qcpUJrSDtH7KR6ZIpB4sijm6ztRlZ4myAe2P0dtaE,101875
+returnn/torch/frontend/_backend.py,sha256=zzKN4_NJK3_I7Ehk8VlhhaXQ_jUEx8K73br8C0Q41p0,103081
 returnn/torch/frontend/_rand.py,sha256=1JgIkV2XmpgJD86zXZ-NCAe-QuoP2swr6NaS1oz3Qa8,1830
 returnn/torch/frontend/bridge.py,sha256=c_mVBCBo29sjm8Bhxarv00szwGPgxjwoIqAHOmceGQw,7842
 returnn/torch/frontend/raw_ops.py,sha256=lF0h-KtYYsdaaqQADylVZp9qzPskOOXA4MfmYDyx5IU,296
@@ -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.20250812.202710.dist-info/LICENSE,sha256=ywBD_U2aD4vpuoIgNAsjIGBYydl0tVKll3De0Z8s77c,11041
-returnn-1.20250812.202710.dist-info/METADATA,sha256=rpwjyEcl0A4JGcT_ukgciH0uSWUHPI8REBpNnhAaPjs,5215
-returnn-1.20250812.202710.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
-returnn-1.20250812.202710.dist-info/top_level.txt,sha256=Lsn4WZc5Pbfk0-xDQOgnFCxOoqxL4CyeM3N1TFbJncw,8
-returnn-1.20250812.202710.dist-info/RECORD,,
+returnn-1.20250817.33823.dist-info/LICENSE,sha256=ywBD_U2aD4vpuoIgNAsjIGBYydl0tVKll3De0Z8s77c,11041
+returnn-1.20250817.33823.dist-info/METADATA,sha256=B57ObmyaxFzp_TjsXhRqtVvzxZWq3hS54vw34tZHZsQ,5214
+returnn-1.20250817.33823.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
+returnn-1.20250817.33823.dist-info/top_level.txt,sha256=Lsn4WZc5Pbfk0-xDQOgnFCxOoqxL4CyeM3N1TFbJncw,8
+returnn-1.20250817.33823.dist-info/RECORD,,