returnn 1.20250117.212735__tar.gz → 1.20250119.205709__tar.gz

{returnn-1.20250117.212735/returnn.egg-info → returnn-1.20250119.205709}/PKG-INFO

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

returnn-1.20250119.205709/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20250119.205709'
+long_version = '1.20250119.205709+git.d3c7c3c'

{returnn-1.20250117.212735 → returnn-1.20250119.205709}/returnn/datasets/postprocessing.py

@@ -88,6 +88,7 @@ class PostprocessingDataset(CachedDataset2):
         map_seq: Optional[Callable] = None,
         map_seq_stream: Optional[Callable] = None,
         map_outputs: Optional[Dict[str, Any]] = None,
+        map_seq_stream_preserves_num_seqs: Optional[bool] = None,
         **kwargs,
     ):
         """
@@ -111,6 +112,8 @@ class PostprocessingDataset(CachedDataset2):
             To simplify the common case when no shapes change, this value can be left unspecified. The dataset then
             assumes the same data layout as returned by the wrapped dataset.
             Example: `map_outputs={"data": {"dim": 42}}`
+        :param map_seq_stream_preserves_num_seqs: whether the function in map_seq_stream preserves the number of
+            sequences, i.e. for every input sequence there is exactly one output sequence.
         :param kwargs: see :class:`CachedDataset2`, :class:`Dataset`
         """
         super().__init__(**kwargs)
@@ -121,19 +124,25 @@ class PostprocessingDataset(CachedDataset2):
             raise ValueError(f"{self}: need to either set map_seq or map_seq_stream")
         if map_seq and map_seq_stream:
             raise ValueError(f"{self}: cannot set both map_seq and map_seq_stream")
+        if map_seq and map_seq_stream_preserves_num_seqs is not None:
+            raise ValueError(f"{self}: map_seq_stream_preserves_num_seqs is only allowed with map_seq_stream")
 
         self._dataset_def = dataset
         self._map_seq = map_seq
         self._map_seq_stream = map_seq_stream
+        if map_seq_stream_preserves_num_seqs is None and map_seq_stream is not None:
+            map_seq_stream_preserves_num_seqs = getattr(map_seq_stream, "preserves_num_seqs", None)
+        self._map_seq_stream_preserves_num_seqs = map_seq_stream_preserves_num_seqs or False
         self._map_outputs = map_outputs
         self._rng = RandomState(self._get_random_seed_for_epoch(0))
         self._seq_list_for_validation: Optional[List[str]] = None
 
         self._dataset = init_dataset(self._dataset_def, parent_dataset=self)
-        if self._map_seq_stream is None:
+        if self._map_seq_stream is None or self._map_seq_stream_preserves_num_seqs:
             # if the stream mapper is set, the num_seqs may change and the estimation is less accurate
             self._estimated_num_seqs = self._dataset.estimated_num_seqs
         self._data_iter: Optional[Iterator[Tuple[int, TensorDict]]] = None
+        self._data_iter_produced_num_seqs = 0
 
         self._in_tensor_dict_template = TensorDict(
             {name: self._make_tensor_template_from_input(name) for name in self._dataset.get_data_keys()}
@@ -187,10 +196,11 @@ 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())
+        self._data_iter_produced_num_seqs = 0
         self._seq_list_for_validation = seq_list
-        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
+        if self._map_seq_stream is None or self._map_seq_stream_preserves_num_seqs:
+            # If we don't have an iterable mapper (or the user explicitly specifies this),
+            # 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:
@@ -221,7 +231,20 @@ class PostprocessingDataset(CachedDataset2):
         while True:
             try:
                 loaded_seq_idx, tensor_dict = next(self._data_iter)
+                self._data_iter_produced_num_seqs += 1
+                if self._num_seqs is not None:
+                    assert self._data_iter_produced_num_seqs <= self._num_seqs, (
+                        f"{self}: map_seq_stream yielded more seqs ({self._data_iter_produced_num_seqs}) "
+                        f"than expected ({self._num_seqs}). map_seq_stream_preserves_num_seqs is set to "
+                        f"{self._map_seq_stream_preserves_num_seqs}"
+                    )
             except StopIteration:
+                if self._num_seqs is not None:
+                    assert self._data_iter_produced_num_seqs == self._num_seqs, (
+                        f"{self}: map_seq_stream yielded {self._data_iter_produced_num_seqs} seqs, "
+                        f"while {self._num_seqs} were expected. map_seq_stream_preserves_num_seqs is set to "
+                        f"{self._map_seq_stream_preserves_num_seqs}"
+                    )
                 return None
             assert loaded_seq_idx <= seq_idx, "_collect_single_seq must be done monotonically"
             if loaded_seq_idx != seq_idx:
@@ -324,6 +347,8 @@ class LaplaceOrdering(Callable[[Iterator[TensorDict]], Iterator[TensorDict]]):
     To be composed with any custom data postprocessing logic via :class:`Sequential`.
     """
 
+    preserves_num_seqs = True
+
     def __init__(self, num_seqs_per_bin: int, length_key: str = "data"):
         """
         :param num_seqs_per_bin: number of segments in a single laplace bin.
@@ -394,3 +419,8 @@ class Sequential:
         for func in self.funcs:
             arg = func(arg, **kwargs)
         return arg
+
+    @property
+    def preserves_num_seqs(self):
+        """:return: whether the composed functions all preserve the number of sequences"""
+        return all(getattr(f, "preserves_num_seqs", False) for f in self.funcs)

{returnn-1.20250117.212735 → returnn-1.20250119.205709}/returnn/torch/frontend/_backend.py

@@ -674,6 +674,12 @@ class TorchBackend(Backend[torch.Tensor]):
         assert targets.sparse_dim and targets.sparse_dim.dimension <= logits.feature_dim.dimension
         # PyTorch expects the logits to be of shape (T, B, C) where T is the input spatial dim.
         batch_dims = logits.remaining_dims((input_spatial_dim, logits.feature_dim))
+        batch_dims_targets = targets.remaining_dims(targets_spatial_dim)
+        if set(batch_dims) != set(batch_dims_targets):
+            # Need to broadcast.
+            logits = rf.expand_dims(logits, [d for d in batch_dims_targets if d not in batch_dims])
+            targets = rf.expand_dims(targets, [d for d in batch_dims if d not in batch_dims_targets])
+            batch_dims = logits.remaining_dims((input_spatial_dim, logits.feature_dim))
         batch_shape = [d.get_dim_value() for d in batch_dims]
         batch_n_elems = prod(batch_shape)
         logits = logits.copy_transpose([input_spatial_dim] + batch_dims + [logits.feature_dim])

{returnn-1.20250117.212735 → returnn-1.20250119.205709/returnn.egg-info}/PKG-INFO

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

returnn-1.20250117.212735/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20250117.212735'
-long_version = '1.20250117.212735+git.7087556'