returnn 1.20251020.105935__py3-none-any.whl → 1.20251020.183915__py3-none-any.whl

returnn/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20251020.105935
+Version: 1.20251020.183915
 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.20251020.105935'
-long_version = '1.20251020.105935+git.d9d5b9f'
+version = '1.20251020.183915'
+long_version = '1.20251020.183915+git.c792649'

returnn/datasets/distrib_files.py

@@ -13,7 +13,7 @@ import sys
 import numpy
 from returnn.log import log
 from returnn.util import better_exchook
-from returnn.util.basic import override_env_var, try_run
+from returnn.util.basic import override_env_var, try_run, OptionalNotImplementedError
 from returnn.util.literal_py_to_pickle import literal_eval
 from returnn.util.multi_proc_non_daemonic_spawn import NonDaemonicSpawnContext
 from returnn.config import SubProcCopyGlobalConfigPreInitFunc
@@ -505,6 +505,24 @@ class DistributeFilesDataset(CachedDataset2):
             self._lazy_init_num_outputs()
         return self._data_keys
 
+    def get_all_tags(self) -> List[str]:
+        """get all tags"""
+        if self.partition_epoch > 1:
+            raise OptionalNotImplementedError(f"{self} get_all_tags not supported for partition_epoch > 1")
+        if self.epoch is None:
+            # Need to init the worker.
+            self.init_seq_order(epoch=1)
+        return self._workers[self.epoch].get_all_tags()
+
+    def get_total_num_seqs(self, *, fast: bool = False) -> int:
+        """get total num seqs"""
+        if self.partition_epoch > 1:
+            raise OptionalNotImplementedError(f"{self} get_total_num_seqs not supported for partition_epoch > 1")
+        if self.epoch is None:
+            # Need to init the worker.
+            self.init_seq_order(epoch=1)
+        return self._workers[self.epoch].get_total_num_seqs(fast=fast)
+
 
 def _get_key_for_file_tree(t: FileTree) -> str:
     """generates a deterministic key given a file tree"""
@@ -608,6 +626,26 @@ class _WorkerProcParent:
         assert msg == "data_seq"
         return data
 
+    def get_all_tags(self) -> List[str]:
+        """get all tags"""
+        self._lazy_wait_for_init_seq_order()
+        self.parent_conn.send(("get_all_tags", {}))
+        msg, data = self.parent_conn.recv()
+        assert msg == "all_tags"
+        if isinstance(data, Exception):
+            raise data
+        return data
+
+    def get_total_num_seqs(self, **kwargs) -> int:
+        """get total num seqs"""
+        self._lazy_wait_for_init_seq_order()
+        self.parent_conn.send(("get_total_num_seqs", kwargs))
+        msg, data = self.parent_conn.recv()
+        assert msg == "total_num_seqs"
+        if isinstance(data, Exception):
+            raise data
+        return data
+
     def exit(self, *, join: bool = True):
         """exit"""
         self._lazy_wait_for_init_seq_order()
@@ -722,6 +760,20 @@ def _worker_proc_loop(
                 got_init_seq_order = True
                 next_seq_idx = 0
                 cache.clear()
+            elif msg == "get_all_tags":
+                try:
+                    tags = dataset.get_all_tags()
+                except Exception as exc:
+                    parent_conn.send(("all_tags", exc))
+                else:
+                    parent_conn.send(("all_tags", tags))
+            elif msg == "get_total_num_seqs":
+                try:
+                    total_num_seqs = dataset.get_total_num_seqs(**kwargs)
+                except Exception as exc:
+                    parent_conn.send(("total_num_seqs", exc))
+                else:
+                    parent_conn.send(("total_num_seqs", total_num_seqs))
             else:
                 raise Exception(f"unknown msg {msg!r}")
     except KeyboardInterrupt:  # when parent dies

returnn/tensor/_dim_extra.py

@@ -163,6 +163,9 @@ class _DimMixin:
     def _handle_extra_kwargs(self: Dim, *, dyn_size: Optional[_t.RawTensorType] = None, **kwargs):
         if kwargs:
             self._extra = _DimExtra(dim=self, **kwargs)
+            if self._extra.vocab and self.size is None:
+                self.size = self._extra.vocab.num_labels
+                self.capacity = self.capacity or self.size
         if dyn_size is not None:
             self.dyn_size = dyn_size
         if self.derived_from_op and self.is_dynamic():

{returnn-1.20251020.105935.dist-info → returnn-1.20251020.183915.dist-info}/METADATA

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

{returnn-1.20251020.105935.dist-info → returnn-1.20251020.183915.dist-info}/RECORD

@@ -1,9 +1,9 @@
-returnn/PKG-INFO,sha256=d-HPlXh1xUoY5v04CS8o4TAtIlgaext6tyRN2PG_lkQ,5215
+returnn/PKG-INFO,sha256=UejLaBSHPrEE1Rzl5GBNrGmToXWMCARlahNqL5lfw5w,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=TchYaay6OjM0dabh0WfmCarERabRL3tGHUI6Squlpzc,77
+returnn/_setup_info_generated.py,sha256=c_EVbbxLj25LjfvHwR99DQ8YX8ShXT13fTsRLyxMz48,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,7 +17,7 @@ returnn/datasets/basic.py,sha256=s0Vjag5lJ5wGXKENN4KHwGtx7ZDiLdWAFIjFbiqAQsE,741
 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=srTieLP02kCepAwZ6Y9p20cqB8nAlVJWbSAoOPna9ik,30567
+returnn/datasets/distrib_files.py,sha256=48edqdf7YpnPJ-TOis3Mz5U9A2DSxfiYT1HCMSti3zw,32718
 returnn/datasets/generating.py,sha256=Qb7V94N_GfL2pZPxWS5PmzszoVXXKzuUmsHuW3dmVbc,99556
 returnn/datasets/hdf.py,sha256=v5sjBenURR9Z-g7AQ9tsL84yDSye5RtbLpym3M6HSDE,67833
 returnn/datasets/huggingface.py,sha256=ls9WMR6gUcMgGksl80g0An1az5Xjya_V3ojbbbsZqrU,20047
@@ -155,7 +155,7 @@ returnn/sprint/extern_interface.py,sha256=l-v1X-Yg0UpTFe7Y3c4FwWOqpSNuv9Oy5EzqlK
 returnn/sprint/interface.py,sha256=1j5SB0V8hSW8A5song9ciZtcBnZoKKfNipk9ezOIMuA,36491
 returnn/tensor/README.md,sha256=X6BqcRLrPLPnwF9yR69uqIFrMnNluj9pBkOPHwNgzuo,501
 returnn/tensor/__init__.py,sha256=on6j5PEOQpck50UcsR4nJzJSDmoVy34z1Oq4efv6Ax0,154
-returnn/tensor/_dim_extra.py,sha256=sZMW5kwyOWOjUdRos5MIxUZOqgBHF99bu9zX2a4HV0Q,116604
+returnn/tensor/_dim_extra.py,sha256=tHE3N6hUKqbzedJ8RNhn9aJHxvhTQuI9JckCLsPbKKI,116776
 returnn/tensor/_tensor_extra.py,sha256=1UPNisRAbljkvfMcrEXaPAF-2Dz7AdgC3jAKVVAnAO8,165084
 returnn/tensor/_tensor_mixin_base.py,sha256=H5z86I0NejxrSgMH1c5oXQzBqS6L9HpvP4y7oegBaSc,643
 returnn/tensor/_tensor_op_overloads.py,sha256=HklwuTBjy7mH_665VKaCUdu-oC3aa7Uz1ZQiCz4jeZc,5448
@@ -254,8 +254,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.20251020.105935.dist-info/LICENSE,sha256=ywBD_U2aD4vpuoIgNAsjIGBYydl0tVKll3De0Z8s77c,11041
-returnn-1.20251020.105935.dist-info/METADATA,sha256=d-HPlXh1xUoY5v04CS8o4TAtIlgaext6tyRN2PG_lkQ,5215
-returnn-1.20251020.105935.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
-returnn-1.20251020.105935.dist-info/top_level.txt,sha256=Lsn4WZc5Pbfk0-xDQOgnFCxOoqxL4CyeM3N1TFbJncw,8
-returnn-1.20251020.105935.dist-info/RECORD,,
+returnn-1.20251020.183915.dist-info/LICENSE,sha256=ywBD_U2aD4vpuoIgNAsjIGBYydl0tVKll3De0Z8s77c,11041
+returnn-1.20251020.183915.dist-info/METADATA,sha256=UejLaBSHPrEE1Rzl5GBNrGmToXWMCARlahNqL5lfw5w,5215
+returnn-1.20251020.183915.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
+returnn-1.20251020.183915.dist-info/top_level.txt,sha256=Lsn4WZc5Pbfk0-xDQOgnFCxOoqxL4CyeM3N1TFbJncw,8
+returnn-1.20251020.183915.dist-info/RECORD,,