returnn 1.20250922.113348__tar.gz → 1.20251006.114241__tar.gz

{returnn-1.20250922.113348/returnn.egg-info → returnn-1.20251006.114241}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20250922.113348
+Version: 1.20251006.114241
 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/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20251006.114241'
+long_version = '1.20251006.114241+git.7745ba7'

{returnn-1.20250922.113348 → returnn-1.20251006.114241}/returnn/datasets/generating.py

@@ -7,14 +7,13 @@ from __future__ import annotations
 from typing import Optional, Union, Any, Sequence, List, Tuple, Dict
 import numpy
 import sys
-import typing
 
 from returnn.util.basic import class_idx_seq_to_1_of_k, CollectionReadCheckCovered
 from returnn.log import log
 from returnn.tensor import Tensor, Dim, TensorDict
 
 from .util.feature_extraction import ExtractAudioFeatures
-from .util.vocabulary import *
+from .util.vocabulary import Vocabulary, BytePairEncoding, CharacterTargets
 from .audio import OggZipDataset  # noqa # for API compatibility
 from .basic import Dataset, DatasetSeq, convert_data_dims
 from .cached2 import CachedDataset2
@@ -1280,12 +1279,6 @@ class CopyTaskDataset(GeneratingDataset):
         return DatasetSeq(seq_idx=seq_idx, features=seq_np, targets={"classes": seq_np})
 
 
-# Multiple external sources where we could write automatic wrappers:
-# * https://github.com/tensorflow/datasets
-# * tf.contrib.keras.datasets, https://www.tensorflow.org/api_docs/python/tf/keras/datasets
-# * nltk.corpus
-
-
 class TimitDataset(CachedDataset2):
     """
     DARPA TIMIT Acoustic-Phonetic Continuous Speech Corpus.
@@ -1553,7 +1546,7 @@ class TimitDataset(CachedDataset2):
 
         self._random_permute_audio = CollectionReadCheckCovered.from_bool_or_dict(random_permute_audio)
 
-        self._seq_order = None  # type: typing.Optional[typing.Sequence[int]]
+        self._seq_order: Optional[Sequence[int]] = None
         self._init_timit()
 
         self._audio_data = {}  # seq_tag -> (audio, sample_rate). loaded by self._reader_thread_main
@@ -1927,8 +1920,8 @@ class BlissDataset(CachedDataset2):
         self._with_delta = with_delta
         self.num_inputs *= 1 + with_delta
         self._bpe_file = open(bpe_file, "r")
-        self._seqs = []  # type: typing.List[BlissDataset.SeqInfo]
-        self._vocab = {}  # type: typing.Dict[str,int]  # set in self._parse_vocab
+        self._seqs: List[BlissDataset.SeqInfo] = []
+        self._vocab: Dict[str, int] = {}  # set in self._parse_vocab
         self._parse_bliss_xml(filename=path)
         # TODO: loading audio like in TimitDataset, and in parallel
         self._bpe = BytePairEncoding(vocab_file=vocab_file, bpe_file=bpe_file)
@@ -2100,7 +2093,7 @@ class LibriSpeechCorpus(CachedDataset2):
             self.targets = CharacterTargets(**chars)
         elif targets is None:
             assert bpe is None and chars is None
-            self.targets = None  # type: typing.Optional[Vocabulary]
+            self.targets: Optional[Vocabulary] = None
         else:
             raise Exception("invalid targets %r. provide bpe or chars" % targets)
         if self.targets:
@@ -2128,7 +2121,7 @@ class LibriSpeechCorpus(CachedDataset2):
             self._reference_seq_order = seqs
             self.transs = {s: self.transs[s] for s in seqs}
         self.epoch_wise_filter = epoch_wise_filter
-        self._seq_order = None  # type: typing.Optional[typing.Sequence[int]]
+        self._seq_order: Optional[Sequence[int]] = None
         self.init_seq_order()
 
     def _collect_trans(self):
@@ -2294,9 +2287,9 @@ class LibriSpeechCorpus(CachedDataset2):
         """:return: whether this dataset supports sharding"""
         return True
 
-    def get_current_seq_order(self):
+    def get_current_seq_order(self) -> Sequence[int]:
         """
-        :rtype: typing.Sequence[int]
+        :return: seq order of current epoch
         """
         assert self._seq_order is not None
         return self._seq_order
@@ -2446,7 +2439,7 @@ class Enwik8Corpus(CachedDataset2):
         self._batch_num_seqs = batch_num_seqs
         self._random = numpy.random.RandomState(1)  # seed will be set in init_seq_order
         self._seq_starts = numpy.arange(0, len(self._data) - 1, seq_len)
-        self._seq_order = None  # type: typing.Optional[typing.Sequence[int]]
+        self._seq_order: Optional[Sequence[int]] = None
 
     def get_data_dtype(self, key):
         """

{returnn-1.20250922.113348 → returnn-1.20251006.114241}/returnn/frontend/_backend.py

@@ -66,6 +66,13 @@ class Backend(Generic[T]):
         """
         raise NotImplementedError
 
+    @staticmethod
+    def should_pickle_tensor(raw_tensor: T) -> bool:
+        """
+        :return: whether the tensor should be included in a pickle or set to `None`.
+        """
+        return True
+
     @staticmethod
     def cond(pred: Tensor, true_fn: Callable, false_fn: Callable):
         """

{returnn-1.20250922.113348 → returnn-1.20251006.114241}/returnn/tensor/_tensor_extra.py

@@ -588,7 +588,12 @@ class _TensorMixin(_TensorMixinBase):
 
     def __getstate__(self):
         d = {k: getattr(self, k) for k in self.__slots__}
-        d["_raw_tensor"] = None  # do not store the TF tensors
+        if (
+            self._raw_tensor is not None
+            and self._raw_backend is not None
+            and not self._raw_backend.should_pickle_tensor(self._raw_tensor)
+        ):
+            d["_raw_tensor"] = None
         return d
 
     def __setstate__(self, state):

{returnn-1.20250922.113348 → returnn-1.20251006.114241}/returnn/tf/frontend_layers/_backend.py

@@ -45,6 +45,13 @@ class ReturnnLayersBackend(Backend[Layer]):
         """executing eagerly"""
         return False
 
+    @staticmethod
+    def should_pickle_tensor(raw_tensor: Layer) -> bool:
+        """
+        :return: whether the tensor should be included in a pickle or set to `None`.
+        """
+        return False
+
     @staticmethod
     def get_tensor_dependencies(x: Tensor[Layer]) -> Sequence[Tensor]:
         """get tensor inputs"""

{returnn-1.20250922.113348 → returnn-1.20251006.114241}/returnn/tf/frontend_low_level/_backend.py

@@ -38,6 +38,21 @@ class TFBackend(Backend[tf.Tensor]):
         """
         return tf.executing_eagerly()
 
+    @staticmethod
+    def should_pickle_tensor(raw_tensor: tf.Tensor) -> bool:
+        """
+        :return: whether the tensor should be included in a pickle or set to `None`.
+        """
+
+        from tensorflow.python.framework.ops import EagerTensor
+
+        # Can not pickle symbolic TF tensors.
+        #
+        # See for discussion:
+        #  - https://github.com/rwth-i6/returnn/issues/1541
+        #  - https://github.com/rwth-i6/returnn/issues/1763
+        return isinstance(raw_tensor, EagerTensor)
+
     @staticmethod
     def get_dtype_name_raw(raw_tensor: tf.Tensor) -> str:
         """

{returnn-1.20250922.113348 → returnn-1.20251006.114241/returnn.egg-info}/PKG-INFO

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

{returnn-1.20250922.113348 → returnn-1.20251006.114241}/tests/test_tensor.py

@@ -2,12 +2,13 @@
 tests for returnn.tensor
 """
 
+import numpy
 import _setup_test_env  # noqa
 import sys
 import unittest
 from returnn.util import better_exchook
 import returnn.frontend as rf
-from returnn.tensor import Tensor, Dim
+from returnn.tensor import Dim, Tensor, TensorDict
 
 
 def test_tensor():
@@ -119,6 +120,52 @@ def test_dim_math_pad_window():
     assert sizes == expected_sizes
 
 
+def test_Tensor_pickle():
+    import pickle
+
+    tensor = Tensor("data", dims=[Dim(10)], dtype="int32", raw_tensor=numpy.zeros((10,), dtype="int32"))
+
+    s = pickle.dumps(tensor)
+    tensor2: Tensor = pickle.loads(s)
+
+    assert tensor.name == tensor2.name
+    assert tensor.dtype == tensor2.dtype
+    assert tensor.raw_tensor is not None
+    assert tensor2.raw_tensor is not None
+    assert numpy.array_equal(tensor.raw_tensor, tensor2.raw_tensor)
+
+
+def test_TensorDict_pickle():
+    import pickle
+
+    tensor_dict = TensorDict(
+        {"data": Tensor("data", dims=[Dim(10)], dtype="int32", raw_tensor=numpy.zeros((10,), dtype="int32"))}
+    )
+    assert all(tensor.raw_tensor is not None for tensor in tensor_dict.data.values())
+
+    s = pickle.dumps(tensor_dict)
+    tensor_dict2 = pickle.loads(s)
+
+    assert tensor_dict.data.keys() == tensor_dict2.data.keys()
+    assert all(tensor.raw_tensor is not None for tensor in tensor_dict2.data.values())
+
+
+def test_TensorDict_queue():
+    from multiprocessing import Queue
+
+    tensor_dict = TensorDict(
+        {"data": Tensor("data", dims=[Dim(10)], dtype="int32", raw_tensor=numpy.zeros((10,), dtype="int32"))}
+    )
+    assert all(tensor.raw_tensor is not None for tensor in tensor_dict.data.values())
+
+    q = Queue(maxsize=2)
+    q.put(tensor_dict)
+    tensor_dict2 = q.get()
+
+    assert tensor_dict.data.keys() == tensor_dict2.data.keys()
+    assert all(tensor.raw_tensor is not None for tensor in tensor_dict2.data.values())
+
+
 if __name__ == "__main__":
     better_exchook.install()
     if len(sys.argv) <= 1:

returnn-1.20250922.113348/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20250922.113348'
-long_version = '1.20250922.113348+git.58aca86'