returnn 1.20250909.170507__tar.gz → 1.20250922.155846__tar.gz

{returnn-1.20250909.170507/returnn.egg-info → returnn-1.20250922.155846}/PKG-INFO

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

@@ -0,0 +1,2 @@
+version = '1.20250922.155846'
+long_version = '1.20250922.155846+git.2334ec6'

{returnn-1.20250909.170507 → returnn-1.20250922.155846}/returnn/datasets/distrib_files.py

@@ -153,7 +153,7 @@ class DistributeFilesDataset(CachedDataset2):
             or a JSON file containing a list of arbitarily nested (JSON) objects.
         :param get_sub_epoch_dataset: callable which returns a dataset dict for a given subset of files
         :param preload_next_n_sub_epochs: how many sub epoch datasets to preload
-        :param buffer_size: buffer size for each worker, amount of seqs to prefetch
+        :param buffer_size: buffer size for each worker, number of seqs to prefetch
         :param distrib_shard_files: set to true to shard the data across worker processes in
             distributed training scenaria
         :param _meta_info_cache: for internal use

{returnn-1.20250909.170507 → returnn-1.20250922.155846}/returnn/datasets/multi_proc.py

@@ -42,7 +42,7 @@ class MultiProcDataset(CachedDataset2):
         """
         :param dataset: the dataset to use
         :param num_workers: number of workers to use
-        :param buffer_size: buffer size for each worker, amount of seqs to prefetch
+        :param buffer_size: buffer size for each worker, number of seqs to prefetch
         :param sharding_method: which method to use for sharding the data across the worker procs.
             The default is ``seq_order``, which fetches the full list of seq indices,
             and then distributes shards of that to the other workers.

{returnn-1.20250909.170507 → returnn-1.20250922.155846}/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.20250909.170507 → returnn-1.20250922.155846}/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.20250909.170507 → returnn-1.20250922.155846}/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.20250909.170507 → returnn-1.20250922.155846}/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.20250909.170507 → returnn-1.20250922.155846/returnn.egg-info}/PKG-INFO

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

{returnn-1.20250909.170507 → returnn-1.20250922.155846}/tests/test_rf_encoder_conformer.py

@@ -496,3 +496,25 @@ def test_e_branchformer_meta_dev_num_params():
         num_params_ = sum(parameter.numel() for parameter in e_branchformer_pt.parameters())
         print(f"E-Branchformer {num_layers}x{dim_} net params #: {num_params_}")
         assert num_params_ == 413034496
+
+        e_branchformer = ConformerEncoder(
+            in_dim=in_dim,
+            input_layer=None,
+            num_layers=num_layers,
+            out_dim=dim,
+            encoder_layer=rf.build_dict(
+                EBranchformerLayer,
+                ff=rf.build_dict(
+                    rf.encoder.conformer.ConformerPositionwiseFeedForward,
+                    ff_dim=dim * 2,
+                    activation=rf.build_dict(rf.relu_square),
+                    with_bias=False,
+                ),
+                num_heads=8,
+            ),
+        )
+        e_branchformer_pt = rf_module_to_pt_module(e_branchformer)
+        assert all(parameter.device.type == "meta" for parameter in e_branchformer_pt.parameters())
+        num_params_ = sum(parameter.numel() for parameter in e_branchformer_pt.parameters())
+        print(f"E-Branchformer {num_layers}x{dim} (FF: {dim}*2) net params #: {num_params_}")
+        assert num_params_ == 404930560

{returnn-1.20250909.170507 → returnn-1.20250922.155846}/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.20250909.170507/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20250909.170507'
-long_version = '1.20250909.170507+git.4dc7088'