returnn 1.20251020.105935__tar.gz → 1.20251020.183915__tar.gz

{returnn-1.20251020.105935/returnn.egg-info → returnn-1.20251020.183915}/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-1.20251020.183915/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20251020.183915'
+long_version = '1.20251020.183915+git.c792649'

{returnn-1.20251020.105935 → returnn-1.20251020.183915}/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-1.20251020.105935 → returnn-1.20251020.183915}/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 → returnn-1.20251020.183915/returnn.egg-info}/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-1.20251020.105935/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20251020.105935'
-long_version = '1.20251020.105935+git.d9d5b9f'