returnn 1.20240908.14935__tar.gz → 1.20240917.154138__tar.gz

{returnn-1.20240908.14935 → returnn-1.20240917.154138}/PKG-INFO

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

returnn-1.20240917.154138/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240917.154138'
+long_version = '1.20240917.154138+git.627ab49'

{returnn-1.20240908.14935 → returnn-1.20240917.154138}/returnn/datasets/postprocessing.py

@@ -9,6 +9,7 @@ from numpy.random import RandomState
 from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, TypeVar
 
 from returnn.datasets.basic import DatasetSeq
+from returnn.datasets.util.strings import str_to_numpy_array
 from returnn.datasets.util.vocabulary import Vocabulary
 from returnn.tensor import Tensor, TensorDict
 from returnn.tensor.dim import Dim
@@ -139,7 +140,9 @@ class PostprocessingDataset(CachedDataset2):
             self._out_tensor_dict_template = TensorDict()
             self._out_tensor_dict_template.update(self._map_outputs, auto_convert=True)
         else:
-            self._out_tensor_dict_template = self._in_tensor_dict_template
+            self._out_tensor_dict_template = self._in_tensor_dict_template.copy_template()
+        # update only after _out_tensor_dict_template has been created from _in_tensor_dict_template
+        self._in_tensor_dict_template.update({"seq_tag": {"dims": (), "dtype": "string"}}, auto_convert=True)
         self.num_outputs = {
             k: (t.sparse_dim.size if t.sparse_dim else t.shape[-1] if len(t.shape) > 0 else 1, t.ndim)
             for k, t in self._out_tensor_dict_template.data.items()
@@ -199,7 +202,11 @@ class PostprocessingDataset(CachedDataset2):
             assert loaded_seq_idx <= seq_idx, "_collect_single_seq must be done monotonically"
             if loaded_seq_idx != seq_idx:
                 continue
-            seq = DatasetSeq(features={k: t.raw_tensor for k, t in tensor_dict.data.items()}, seq_idx=seq_idx)
+            seq = DatasetSeq(
+                features={k: t.raw_tensor for k, t in tensor_dict.data.items() if k != "seq_tag"},
+                seq_idx=seq_idx,
+                seq_tag=str(tensor_dict["seq_tag"].raw_tensor),
+            )
             return seq
 
     def _build_mapping_iter(self) -> Iterator[TensorDict]:
@@ -209,6 +216,7 @@ class PostprocessingDataset(CachedDataset2):
 
         def _validate_tensor_dict_iter(inner: Iterator[TensorDict]) -> Iterator[TensorDict]:
             for t_dict in inner:
+                assert "seq_tag" in t_dict.data, "seq_tag dropped from TensorDict in postprocessing pipeline"
                 for data_key, out_t in self._out_tensor_dict_template.data.items():
                     in_t = t_dict.data[data_key]
                     assert (
@@ -237,9 +245,12 @@ class PostprocessingDataset(CachedDataset2):
         seq_index = 0
         while self._dataset.is_less_than_num_seqs(seq_index):
             self._dataset.load_seqs(seq_index, seq_index + 1)
+
             tensor_dict = self._in_tensor_dict_template.copy_template()
             for data_key in data_keys:
                 tensor_dict.data[data_key].raw_tensor = self._dataset.get_data(seq_index, data_key)
+            tensor_dict.data["seq_tag"].raw_tensor = str_to_numpy_array(self._dataset.get_tag(seq_index))
+
             if self._map_seq is not None:
                 tensor_dict = self._map_seq(
                     tensor_dict, rng=self._rng, **{f"fwd_compatible_random_kwarg_{self._rng.randint(0, 1000)}": None}
@@ -247,6 +258,12 @@ class PostprocessingDataset(CachedDataset2):
                 assert isinstance(
                     tensor_dict, TensorDict
                 ), f"map_seq must produce a {TensorDict.__name__}, but produced {type(tensor_dict).__name__}"
+
+                # Re-adding the seq tag here causes no harm in case it's dropped since we don't
+                # add/drop any segments w/ the non-iterator postprocessing function.
+                if "seq_tag" not in tensor_dict.data:
+                    tensor_dict.data["seq_tag"].raw_tensor = str_to_numpy_array(self._dataset.get_tag(seq_index))
+
             yield tensor_dict
             seq_index += 1
 

{returnn-1.20240908.14935 → returnn-1.20240917.154138}/returnn/frontend/run_ctx.py

@@ -9,6 +9,7 @@ or forwarding loop.
 from __future__ import annotations
 from typing import Optional, Union, Any, Sequence, Dict
 from dataclasses import dataclass
+from contextlib import contextmanager
 from returnn.tensor import Tensor, Dim, TensorDict, batch_dim
 import returnn.frontend as rf
 from . import _backend
@@ -126,6 +127,25 @@ class RunCtx:
         """
         return self._train_flag
 
+    @contextmanager
+    def train_flag_ctx(self, train_flag: Union[bool, Tensor]):
+        """
+        Context manager to temporarily set the train_flag.
+
+        Usage example e.g. to disable dropout for some code::
+
+            with rf.get_run_ctx().train_flag_ctx(False):
+                ...
+
+        :param train_flag: whether we are in training mode
+        """
+        old_train_flag = self.train_flag
+        self._train_flag = train_flag
+        try:
+            yield
+        finally:
+            self._train_flag = old_train_flag
+
     @property
     def step(self) -> Union[int, Tensor]:
         """

{returnn-1.20240908.14935 → returnn-1.20240917.154138}/returnn/tensor/_dim_extra.py

@@ -136,6 +136,9 @@ class _DimExtra:
         if self.kind is not None:
             self.kind = {v.name: v for v in DimTypes.Types}[self.kind]
 
+    def __sis_state__(self):
+        raise ValueError(f"{self}: currently not expected to be part of the Sisyphus state/hash")
+
 
 class _DimMixin:
     name: Optional[str]
@@ -2280,7 +2283,7 @@ class _DimMixin:
 
     def _cache_dim_math_get(
         self: Dim, op_kind: str, operand: Union[Dim, int]
-    ) -> Tuple[Dict[Tuple[str, Union[Dim, int]], Dim], Tuple[str, Union[Dim, int]], Optional[Dim]]:
+    ) -> Tuple[_CacheDimMath, Tuple[str, Union[Dim, int]], Optional[Dim]]:
         same_base = self.get_same_base()
         # noinspection PyProtectedMember
         extra = same_base._make_extra()
@@ -2290,6 +2293,9 @@ class _DimMixin:
         cache_key = (op_kind, operand)
         return cache, cache_key, cache.get(cache_key, None)
 
+    def __sis_state__(self):
+        raise ValueError(f"{self}: currently not expected to be part of the Sisyphus state/hash")
+
 
 def _make_constant_static_dim(value, kind=None):
     """

{returnn-1.20240908.14935 → returnn-1.20240917.154138}/returnn.egg-info/PKG-INFO

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

{returnn-1.20240908.14935 → returnn-1.20240917.154138}/tests/test_Dataset.py

@@ -1086,6 +1086,9 @@ def test_PostprocessingDataset():
         classes = dataset.get_data(0, "classes")
         assert all(c - 1337 >= 0 for c in classes)
 
+        # assert that default seq tags have been replaced w/ ones from oggzip dset
+        assert not dataset.get_tag(0).startswith("seq-")
+
     count = 0
 
     def _repeat2(input_iter: Iterator[TensorDict], **kwargs) -> Iterator[TensorDict]:

returnn-1.20240908.14935/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240908.014935'
-long_version = '1.20240908.014935+git.7390bea'