returnn 1.20240905.141151__tar.gz → 1.20240905.172412__tar.gz

{returnn-1.20240905.141151 → returnn-1.20240905.172412}/PKG-INFO

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

returnn-1.20240905.172412/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240905.172412'
+long_version = '1.20240905.172412+git.fb9d5c3'

{returnn-1.20240905.141151 → returnn-1.20240905.172412}/returnn/datasets/distrib_files.py

@@ -417,8 +417,6 @@ class DistributeFilesDataset(CachedDataset2):
         return files_per_bin
 
     def _collect_single_seq(self, seq_idx: int) -> Optional[DatasetSeq]:
-        if seq_idx >= self._num_seqs:
-            return None
         seq = self._workers[self.epoch].get_data_seq(seq_idx)
         if not seq:
             return None
@@ -429,11 +427,6 @@ class DistributeFilesDataset(CachedDataset2):
         """have seqs"""
         return bool(self.files)
 
-    @property
-    def num_seqs(self) -> int:
-        """num seqs"""
-        return self._num_seqs
-
     def finish_epoch(self, *, free_resources: bool = False):
         """finish epoch"""
         super().finish_epoch(free_resources=free_resources)
@@ -530,16 +523,18 @@ class _WorkerProcParent:
         # Thus, use the full epoch as the epoch here.
         # Make sure this init is async - we should not wait for this in the main process when it is not used yet.
         self.parent_conn.send(("init_seq_order", {"epoch": self.full_epoch_0idx + 1}))
-        self._num_seqs: Optional[int] = None
+        self._has_num_seqs = False
+        self._num_seqs: Optional[int] = None  # remains None if unknown
 
     def _lazy_wait_for_init_seq_order(self):
-        if self._num_seqs is not None:
+        if self._has_num_seqs:
             return
         msg, num_seqs = self.parent_conn.recv()
-        assert msg == "num_seqs" and isinstance(num_seqs, int)
+        assert msg == "num_seqs" and (num_seqs is None or isinstance(num_seqs, int))
         self._num_seqs = num_seqs
+        self._has_num_seqs = True
 
-    def get_num_seqs(self) -> int:
+    def get_num_seqs(self) -> Optional[int]:
         """num seqs for this sub epoch"""
         self._lazy_wait_for_init_seq_order()
         return self._num_seqs
@@ -656,7 +651,11 @@ def _worker_proc_loop(
             elif msg == "init_seq_order":
                 # We are responsible to get the seq order and distrib it to all the other workers.
                 dataset.init_seq_order(**kwargs)
-                parent_conn.send(("num_seqs", dataset.num_seqs))
+                try:
+                    num_seqs = dataset.num_seqs
+                except NotImplementedError:
+                    num_seqs = None
+                parent_conn.send(("num_seqs", num_seqs))
                 got_init_seq_order = True
                 next_seq_idx = 0
                 cache[:] = []

{returnn-1.20240905.141151 → returnn-1.20240905.172412}/returnn/datasets/postprocessing.py

@@ -173,6 +173,13 @@ class PostprocessingDataset(CachedDataset2):
         assert self._dataset is not None
         self._dataset.init_seq_order(epoch=epoch, seq_list=seq_list, seq_order=seq_order)
         self._data_iter = enumerate(self._build_mapping_iter())
+        if self._map_seq_stream is None:
+            # If we don't have an iterable mapper we know the number of segments exactly
+            # equals the number of segments in the wrapped dataset
+            try:
+                self._num_seqs = self._dataset.num_seqs
+            except NotImplementedError:
+                pass  # some datasets don't know their num_seqs
         return True
 
     def get_data_keys(self):

{returnn-1.20240905.141151 → returnn-1.20240905.172412}/returnn/frontend/_backend.py

@@ -957,6 +957,11 @@ class Backend(Generic[T]):
         """slice"""
         raise NotImplementedError
 
+    @staticmethod
+    def flip(source: Tensor, *, axis: Dim) -> Tensor:
+        """flip"""
+        raise NotImplementedError
+
     @staticmethod
     def where(
         cond: Tensor,

{returnn-1.20240905.141151 → returnn-1.20240905.172412}/returnn/frontend/array_.py

@@ -723,11 +723,14 @@ def slice(
     :param axis:
     :param start:
     :param end:
-    :param step:
+    :param step: we only support a positive step here. If you want to flip the sequence,
+        use :func:`reverse_sequence`, maybe with ``handle_dynamic_dims=False``.
     :param size:
     :param out_dim:
     :return: tensor, out_dim
     """
+    if isinstance(step, int) and step <= 0:
+        raise ValueError(f"rf.slice with non-positive step {step} not supported")
     if not out_dim:
         # Note: We cover not really all possible cases here but just those we really support.
         # (Actually, it might still be a superset of cases we really support,
@@ -794,14 +797,21 @@ def shift_right(source: Tensor, *, axis: Dim, pad_value: Union[rf.RawTensorTypes
     return padded_slice
 
 
-def reverse_sequence(tensor: Tensor, *, axis: Dim) -> Tensor:
+def reverse_sequence(tensor: Tensor, *, axis: Dim, handle_dynamic_dims: bool = True) -> Tensor:
     """
     Similar as tf.reverse_sequence, or Torch flip (but taking seq lengths into account).
 
     :param tensor:
     :param axis:
+    :param handle_dynamic_dims: if True, the dynamic dims will be handled correctly,
+        i.e. a padded seq "a b c <pad>" will be reversed to "c b a <pad>".
+        If False, the dynamic dims will be handled as if they are static,
+        i.e. a padded seq "a b c <pad>" will be reversed to "<pad> c b a".
     :return: reversed tensor, same dims
     """
+    if not handle_dynamic_dims or not axis.need_masking():
+        # noinspection PyProtectedMember
+        return tensor._raw_backend.flip(tensor, axis=axis)
     indices = rf.combine_bc(axis.get_size_tensor(), "-", rf.range_over_dim(axis)) - 1
     return rf.gather(tensor, indices=indices, axis=axis, clip_to_valid=True)
 

{returnn-1.20240905.141151 → returnn-1.20240905.172412}/returnn/tf/frontend_layers/_backend.py

@@ -695,6 +695,13 @@ class ReturnnLayersBackend(Backend[Layer]):
             {"class": "slice", "from": source, "axis": axis, "out_dim": out_dim, **args}, name="slice"
         )
 
+    @staticmethod
+    def flip(source: Tensor, *, axis: Dim) -> Tensor:
+        """flip"""
+        return rfl.make_layer(
+            {"class": "slice", "from": source, "axis": axis, "out_dim": axis, "slice_step": -1}, name="flip"
+        )
+
     @staticmethod
     def where(
         cond: Tensor,

{returnn-1.20240905.141151 → returnn-1.20240905.172412}/returnn/torch/frontend/_backend.py

@@ -1096,6 +1096,14 @@ class TorchBackend(Backend[torch.Tensor]):
             out.raw_tensor = torch.narrow(source.raw_tensor, dim=axis_int, start=start, length=end - start)
         return out
 
+    @staticmethod
+    def flip(source: Tensor, *, axis: Dim) -> Tensor:
+        """flip"""
+        axis_int = source.get_axis_from_description(axis, allow_int=False)
+        out = source.copy_template("flip")
+        out.raw_tensor = torch.flip(source.raw_tensor, [axis_int])
+        return out
+
     @staticmethod
     def where(
         cond: Tensor,

{returnn-1.20240905.141151 → returnn-1.20240905.172412}/returnn.egg-info/PKG-INFO

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

{returnn-1.20240905.141151 → returnn-1.20240905.172412}/tests/test_Dataset.py

@@ -943,6 +943,10 @@ def test_DistributeFilesDataset_distribute_evenly_by_size():
     _test_stddev(rng.uniform(low=97, high=103, size=(100,)), 20, 0.03)
 
 
+def iter_identity(x, **kwargs):
+    yield from x
+
+
 def test_DistributeFilesDataset():
     from returnn.datasets.distrib_files import DistributeFilesDataset
     from test_HDFDataset import generate_hdf_from_other
@@ -1036,6 +1040,27 @@ def test_DistributeFilesDataset():
             global_seq_idx += 1
     assert global_seq_idx == num_hdf_files * num_seqs
 
+    # Test DistributeFilesDataset in conjunction w/ a dataset, that does not have a
+    # defined num_seqs:
+
+    def _get_subepoch_dset_w_unknown_num_seqs(files_subepoch: List[str]) -> Dict[str, Any]:
+        hdf_dset = {"class": "HDFDataset", "files": files_subepoch, "seq_ordering": "default"}
+        return {"class": "PostprocessingDataset", "dataset": hdf_dset, "map_seq_stream": iter_identity}
+
+    concat_dataset = init_dataset(
+        {
+            "class": "DistributeFilesDataset",
+            "files": hdf_files,
+            "get_sub_epoch_dataset": _get_subepoch_dset_w_unknown_num_seqs,
+            "partition_epoch": partition_epoch,
+        }
+    )
+    assert isinstance(concat_dataset, DistributeFilesDataset)
+    for sub_epoch in range(1, partition_epoch + 1):
+        print(f"Sub-epoch {sub_epoch}...")
+        concat_dataset.init_seq_order(epoch=sub_epoch)
+        concat_dataset.load_seqs(0, 1)
+
 
 def test_PostprocessingDataset():
     from returnn.tensor.tensor_dict import TensorDict

{returnn-1.20240905.141151 → returnn-1.20240905.172412}/tests/test_rf_array.py

@@ -398,6 +398,23 @@ def test_reverse_sequence():
     run_model(extern_data, lambda *, epoch, step: _Net(), _forward_step)
 
 
+def test_reverse_sequence_no_dyn():
+    time_dim = Dim(Tensor("time", [batch_dim], dtype="int32"))
+    in_dim = Dim(7, name="in")
+    extern_data = TensorDict(
+        {
+            "data": Tensor("data", [batch_dim, time_dim, in_dim], dtype="float32"),
+        }
+    )
+
+    # noinspection PyShadowingNames
+    def _forward_step(*, extern_data: TensorDict, **_kwargs):
+        out = rf.reverse_sequence(extern_data["data"], axis=time_dim, handle_dynamic_dims=False)
+        out.mark_as_default_output(shape=(batch_dim, time_dim, in_dim))
+
+    run_model(extern_data, lambda *, epoch, step: rf.Module(), _forward_step)
+
+
 def test_where():
     time_dim = Dim(Tensor("time", [batch_dim], dtype="int32"))
     in_dim = Dim(7, name="in")

returnn-1.20240905.141151/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240905.141151'
-long_version = '1.20240905.141151+git.ce97976'