returnn 1.20250226.132109__tar.gz → 1.20250227.110407__tar.gz

{returnn-1.20250226.132109/returnn.egg-info → returnn-1.20250227.110407}/PKG-INFO

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

returnn-1.20250227.110407/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20250227.110407'
+long_version = '1.20250227.110407+git.f069d06'

{returnn-1.20250226.132109 → returnn-1.20250227.110407}/returnn/datasets/basic.py

@@ -16,6 +16,7 @@ from threading import RLock
 from random import Random, random
 import sys
 import os
+import math
 import numpy
 import functools
 import typing
@@ -937,28 +938,51 @@ class Dataset:
         else:
             # We don't know. So:
             # Some monotonic increasing function in [0,1] which never reaches 1.
-            import math
+            return max(1.0e-10, (1 - 1 / ((seq_idx**0.5) / 100 + 1)) * 0.99)
 
-            return max(1.0e-10, 1.0 - math.exp(-seq_idx * 1000))
-
-    def get_complete_frac(self, seq_idx):
+    def get_complete_frac(self, sorted_seq_idx: int, *, allow_only_lr_suitable: bool = False) -> Optional[float]:
         """
-        :param int seq_idx:
-        :return: Returns a fraction (float in [0,1], always > 0) of how far we have advanced
-          for this seq in the dataset.
-          This does not have to be exact. This is only for the user.
-        :rtype: float
+        Tries to calculate exactly how much of the current epoch is completed when
+        having processed seq ``sorted_seq_idx``.
+
+        ``sorted_seq_idx`` cannot be less than the seq index of the previously loaded seqs.
+
+        :param sorted_seq_idx: sorted seq idx
+        :param allow_only_lr_suitable: only return a value when that value is suitable/accurate enough
+            to base LR scheduling on it. If false, this function will return an approximative value
+            when the exact value cannot be calculated (due to unknown ``num_seqs``).
+            Approximative values can be appropriate for e.g. progress bars.
+        :return: continuous value in (0, 1] which represents how much of the current epoch
+            is completed after ``sorted_seq_idx``.
+            If ``allow_only_lr_suitable=True``, returns ``None`` if the value cannot be calculated such
+            that it is accurate enough for LR scheduling, and otherwises bases ``epoch_continuous`` on it
+            for any dynamic learning rate scheduling.
+            As ``sorted_seq_idx`` is monotonic, the return value is also guaranteed to be monotonic.
         """
         # noinspection PyBroadException
         try:
             num_seqs = self.num_seqs
         except Exception:  # num_seqs not always available
+            if allow_only_lr_suitable:
+                return None
+
             # noinspection PyBroadException
             try:
                 num_seqs = self.estimated_num_seqs
             except Exception:  # also not always available
                 num_seqs = None  # ignore
-        return self.generic_complete_frac(seq_idx, num_seqs)
+
+        if math.isinf(num_seqs):
+            if allow_only_lr_suitable:
+                # cannot compute meaningful complete_frac for infinite num_seqs
+                return None
+            else:
+                num_seqs = None
+
+        assert (
+            num_seqs is None or 0 <= sorted_seq_idx < num_seqs
+        ), f"{self}: invalid seq indices: 0 <= seq_idx ({sorted_seq_idx}) < num_seqs ({num_seqs}) violated"
+        return self.generic_complete_frac(sorted_seq_idx, num_seqs)
 
     @property
     def num_seqs(self) -> int:
@@ -1375,16 +1399,27 @@ class DatasetSeq:
     Encapsulates all data for one sequence.
     """
 
-    def __init__(self, seq_idx, features, targets=None, seq_tag=None):
+    def __init__(
+        self,
+        seq_idx: int,
+        features,
+        *,
+        targets=None,
+        seq_tag: Optional[str] = None,
+        complete_frac: Optional[float] = None,
+    ):
         """
-        :param int seq_idx: sorted seq idx in the Dataset
+        :param seq_idx: sorted seq idx in the Dataset
         :param numpy.ndarray|dict[str,numpy.ndarray] features: format 2d (time,feature) (float)
         :param dict[str,numpy.ndarray]|numpy.ndarray|None targets: name -> format 1d (time) (idx of output-feature)
-        :param str seq_tag: sequence name / tag
+        :param seq_tag: sequence name / tag
+        :param complete_frac: continuous value in (0, 1] which represents how much of the current epoch
+            has been consumed when this seq is processed
         """
         assert isinstance(seq_idx, (int, numpy.integer))
         self.seq_idx = int(seq_idx)
         self.seq_tag = seq_tag or ("seq-%i" % seq_idx)
+        self.complete_frac = complete_frac
         if not isinstance(features, dict):
             assert isinstance(features, numpy.ndarray)
             features = {"data": features}

{returnn-1.20250226.132109 → returnn-1.20250227.110407}/returnn/datasets/cached2.py

@@ -228,6 +228,15 @@ class CachedDataset2(Dataset):
             keys.remove("data")
         return keys
 
+    def get_complete_frac(self, sorted_seq_idx, **kwargs):
+        """
+        :return: fractional completion value for the given sorted_seq_idx
+        """
+        seq = self._get_seq(sorted_seq_idx)
+        if seq is not None and seq.complete_frac is not None:
+            return seq.complete_frac
+        return super().get_complete_frac(sorted_seq_idx, **kwargs)
+
     def is_data_sparse(self, key):
         """
         :param str key: e.g. "data" or "classes"

{returnn-1.20250226.132109 → returnn-1.20250227.110407}/returnn/datasets/distrib_files.py

@@ -601,7 +601,8 @@ def _worker_proc_loop(
         dataset.load_seqs(next_seq_idx, next_seq_idx + 1)
         seq_tag = dataset.get_tag(next_seq_idx)
         features = {data_key: dataset.get_data(next_seq_idx, data_key) for data_key in dataset.get_data_keys()}
-        res = DatasetSeq(seq_idx=next_seq_idx, seq_tag=seq_tag, features=features)
+        complete_frac = dataset.get_complete_frac(next_seq_idx, allow_only_lr_suitable=True)
+        res = DatasetSeq(seq_idx=next_seq_idx, seq_tag=seq_tag, features=features, complete_frac=complete_frac)
         cache.append(res)
         next_seq_idx += 1
         return True

{returnn-1.20250226.132109 → returnn-1.20250227.110407}/returnn/datasets/meta.py

@@ -554,6 +554,12 @@ class MetaDataset(CachedDataset2):
         """
         return self.seq_list_ordered[self.default_dataset_key][sorted_seq_idx]
 
+    def get_complete_frac(self, sorted_seq_idx: int, **kwargs) -> Optional[float]:
+        """
+        :param sorted_seq_idx:
+        """
+        return self.datasets[self.default_dataset_key].get_complete_frac(sorted_seq_idx, **kwargs)
+
     def get_data_keys(self) -> List[str]:
         """data keys"""
         return sorted(self.data_keys)

{returnn-1.20250226.132109 → returnn-1.20250227.110407}/returnn/datasets/multi_proc.py

@@ -75,6 +75,7 @@ class MultiProcDataset(CachedDataset2):
         self._seq_order_proc_parent_conn = None  # type: Optional[mpConnection]
         self._seq_order_proc = None  # type: Optional[mp.Process]
         self._worker_procs = None  # type: Optional[List[mp.Process]]
+        self._cur_max_complete_frac: Optional[float] = None
 
         if _meta_info_cache:
             # This allows to skip the lazy init in self.initialize().
@@ -246,7 +247,8 @@ class MultiProcDataset(CachedDataset2):
             dataset.load_seqs(next_seq_idx, next_seq_idx + 1)
             seq_tag = dataset.get_tag(next_seq_idx)
             features = {data_key: dataset.get_data(next_seq_idx, data_key) for data_key in dataset.get_data_keys()}
-            res = DatasetSeq(seq_idx=next_seq_idx, seq_tag=seq_tag, features=features)
+            complete_frac = dataset.get_complete_frac(next_seq_idx, allow_only_lr_suitable=True)
+            res = DatasetSeq(seq_idx=next_seq_idx, seq_tag=seq_tag, features=features, complete_frac=complete_frac)
             cache.append(res)
             next_seq_idx += 1
             return True
@@ -403,6 +405,7 @@ class MultiProcDataset(CachedDataset2):
             return True
 
         self._lazy_init()
+        self._cur_max_complete_frac = 0.0
 
         if self._sharding_method == "dedicated":
             for worker_conn in self._worker_parent_conns:
@@ -441,6 +444,12 @@ class MultiProcDataset(CachedDataset2):
         if data is None:
             return None
         assert isinstance(data, DatasetSeq)
+        # The complete_frac values from the subprocesses are not necessarily monotonic
+        # due to rounding errors in the sharding and such.
+        # We therefore fix them up here. This is valid due to monotonicity of `seq_idx`.
+        max_comp_frac = max(data.complete_frac, self._cur_max_complete_frac)
+        data.complete_frac = max_comp_frac
+        self._cur_max_complete_frac = max_comp_frac
         data.seq_idx = seq_idx
         return data
 

{returnn-1.20250226.132109 → returnn-1.20250227.110407}/returnn/datasets/numpy_dump.py

@@ -154,4 +154,4 @@ class NumpyDumpDataset(Dataset):
     def _add_cache_seq(self, seq_idx, features, targets):
         last_seq_idx = self._get_cache_last_seq_idx()
         assert seq_idx == last_seq_idx + 1
-        self.cached_seqs += [DatasetSeq(seq_idx, features, targets)]
+        self.cached_seqs += [DatasetSeq(seq_idx, features, targets=targets)]

{returnn-1.20250226.132109 → returnn-1.20250227.110407}/returnn/datasets/postprocessing.py

@@ -5,6 +5,7 @@ Provides :class:`PostprocessingDataset`.
 from __future__ import annotations
 
 from itertools import islice
+import numpy
 from numpy.random import RandomState
 from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, TypeVar
 
@@ -57,6 +58,14 @@ class PostprocessingDataset(CachedDataset2):
             },
         }
 
+    The postprocessor functions operate on ``TensorDict``s, which have entries for
+    all data keys in the underlying dataset.
+
+    There may also be additional "meta" entries in the tensor dicts, like ``complete_frac``
+    and ``seq_tag``.
+    These should be copied over in a manner that is reasonable for the use case at hand and
+    ensures forwards compatibility as well as reasonably possible.
+
     The dataset itself does not support its own seq ordering and relies on the wrapped
     dataset for seq ordering instead. Specifying a ``seq_ordering`` other than ``default``
     results in an error.
@@ -155,7 +164,10 @@ class PostprocessingDataset(CachedDataset2):
             self._out_tensor_dict_template = self._in_tensor_dict_template.copy_template()
             self.labels = self._dataset.labels.copy()
         # 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._in_tensor_dict_template.update(
+            {"complete_frac": {"dims": (), "dtype": "float32"}, "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()
@@ -222,6 +234,15 @@ class PostprocessingDataset(CachedDataset2):
         """:return: dtype of data entry `key`"""
         return self._out_tensor_dict_template.data[key].dtype
 
+    def get_total_num_seqs(self, *, fast=False):
+        """:return: total num seqs excluding partition_epoch"""
+        if self._map_seq_stream is not None:
+            raise util.OptionalNotImplementedError(
+                f"{self}: get_total_num_seqs not allowed when map_seq_stream is set."
+            )
+        assert self._dataset is not None
+        return self._dataset.get_total_num_seqs(fast=fast)
+
     def supports_sharding(self) -> bool:
         """:return: whether this dataset supports sharding"""
         assert self._dataset is not None
@@ -249,11 +270,12 @@ 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() if k != "seq_tag"},
-                seq_idx=seq_idx,
-                seq_tag=str(tensor_dict["seq_tag"].raw_tensor),
+            complete_frac = (
+                float(tensor_dict.data["complete_frac"].raw_tensor) if "complete_frac" in tensor_dict.data else None
             )
+            seq_tag = str(tensor_dict.data["seq_tag"].raw_tensor) if "seq_tag" in tensor_dict.data else f"seq-{seq_idx}"
+            features = {k: t.raw_tensor for k, t in tensor_dict.data.items() if k not in ["complete_frac", "seq_tag"]}
+            seq = DatasetSeq(complete_frac=complete_frac, features=features, seq_idx=seq_idx, seq_tag=seq_tag)
             return seq
 
     def _build_mapping_iter(self) -> Iterator[TensorDict]:
@@ -262,8 +284,20 @@ class PostprocessingDataset(CachedDataset2):
         """
 
         def _validate_tensor_dict_iter(inner: Iterator[TensorDict]) -> Iterator[TensorDict]:
+            last_complete_frac = 0.0
             for t_dict in inner:
-                assert "seq_tag" in t_dict.data, "seq_tag dropped from TensorDict in postprocessing pipeline"
+                assert isinstance(t_dict, TensorDict), (
+                    f"postprocessing mapper function must produce a {TensorDict.__name__}, "
+                    f"but got a {type(t_dict).__name__}"
+                )
+                if "complete_frac" in t_dict.data:  # sanity check complete_frac
+                    complete_frac = float(t_dict.data["complete_frac"].raw_tensor)
+                    assert 0.0 <= complete_frac <= 1.0, f"complete_frac must be in [0, 1], but got {complete_frac}"
+                    assert complete_frac >= last_complete_frac, (
+                        "complete_frac must be monotonically increasing, "
+                        f"but got {complete_frac} after {last_complete_frac}"
+                    )
+                    last_complete_frac = complete_frac
                 for data_key, out_t in self._out_tensor_dict_template.data.items():
                     in_t = t_dict.data[data_key]
                     assert (
@@ -294,8 +328,14 @@ class PostprocessingDataset(CachedDataset2):
             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)
-            seq_tag_tensor = str_to_numpy_array(self._dataset.get_tag(seq_index))
-            tensor_dict.data["seq_tag"].raw_tensor = seq_tag_tensor
+
+            complete_frac = self._dataset.get_complete_frac(seq_index, allow_only_lr_suitable=True)
+            comp_frac_raw_tensor = None
+            if complete_frac is not None:
+                comp_frac_raw_tensor = numpy.array(complete_frac, dtype=numpy.float32)
+                tensor_dict.data["complete_frac"].raw_tensor = comp_frac_raw_tensor
+            seq_tag_raw_tensor = str_to_numpy_array(self._dataset.get_tag(seq_index))
+            tensor_dict.data["seq_tag"].raw_tensor = seq_tag_raw_tensor
 
             if self._map_seq is not None:
                 tensor_dict = self._map_seq(
@@ -305,10 +345,16 @@ class PostprocessingDataset(CachedDataset2):
                     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.
+                # Re-adding the seq_tag/complete_frac here causes no harm in case they are dropped
+                # since we don't add/drop any segments w/ the non-iterator postprocessing function.
+                if "complete_frac" not in tensor_dict.data and comp_frac_raw_tensor is not None:
+                    tensor_dict.data["complete_frac"] = Tensor(
+                        "complete_frac", dims=(), dtype="float32", raw_tensor=comp_frac_raw_tensor
+                    )
                 if "seq_tag" not in tensor_dict.data:
-                    tensor_dict.data["seq_tag"].raw_tensor = seq_tag_tensor
+                    tensor_dict.data["seq_tag"] = Tensor(
+                        "seq_tag", dims=(), dtype="string", raw_tensor=seq_tag_raw_tensor
+                    )
 
                 if self._seq_list_for_validation is not None:
                     seq_tag = self._seq_list_for_validation[seq_index]
@@ -366,7 +412,12 @@ class LaplaceOrdering(Callable[[Iterator[TensorDict]], Iterator[TensorDict]]):
         seq_buffer = list(islice(iterator, self.num_seqs_per_bin))
         has_ended = False
         while True:
+            # Make sure to not reorder the monotonically increasing values for complete_frac
+            # so that the trainer can calculate the appropriate learning rates.
+            complete_frac_values = [tdict.data["complete_frac"].raw_tensor for tdict in seq_buffer]
             seq_buffer.sort(key=self._get_seq_len, reverse=is_down_phase)
+            for sorted_item, comp_frac in zip(seq_buffer, complete_frac_values):
+                sorted_item.data["complete_frac"].raw_tensor = comp_frac
 
             next_seq_buffer = []
 

{returnn-1.20250226.132109 → returnn-1.20250227.110407}/returnn/datasets/raw_wav.py

@@ -83,7 +83,7 @@ class RawWavDataset(CachedDataset2):
         inputFeatures = inputFeatures.astype(np.float32)
         if outputFeatures is not None:
             outputFeatures = outputFeatures.astype(np.float32)
-        return DatasetSeq(seq_idx, inputFeatures, outputFeatures)
+        return DatasetSeq(seq_idx, inputFeatures, targets=outputFeatures)
 
     def _get_num_outputs(self, num_outputs):
         """

{returnn-1.20250226.132109 → returnn-1.20250227.110407}/returnn/datasets/sprint.py

@@ -504,7 +504,7 @@ class SprintDatasetBase(Dataset):
                     assert seq_idx + 1 == self.next_seq_to_be_added
                     self.cond.wait()
 
-            self.added_data += [DatasetSeq(seq_idx, features, targets, seq_tag=segment_name)]
+            self.added_data += [DatasetSeq(seq_idx, features, targets=targets, seq_tag=segment_name)]
             self.cond.notify_all()
             return seq_idx
 
@@ -588,7 +588,7 @@ class SprintDatasetBase(Dataset):
         """
         self._complete_frac = frac
 
-    def get_complete_frac(self, seq_idx):
+    def get_complete_frac(self, seq_idx, **kwargs):
         """
         :param int seq_idx:
         :rtype: float

{returnn-1.20250226.132109 → returnn-1.20250227.110407}/returnn/datasets/stereo.py

@@ -349,7 +349,7 @@ class StereoHdfDataset(StereoDataset):
         elif targets.shape[1] == 1:
             targets = np.reshape(targets.astype(np.int32), (targets.shape[0],))
 
-        return DatasetSeq(seq_idx, inputFeatures, targets)
+        return DatasetSeq(seq_idx, inputFeatures, targets=targets)
 
     @staticmethod
     def _normalizeVector(v, mean, variance):
@@ -438,4 +438,4 @@ class DatasetWithTimeContext(StereoHdfDataset):
         targets = None
         if "classes" in originalSeq.get_data_keys():
             targets = originalSeq.get_data("classes")
-        return DatasetSeq(seq_idx, inputFeatures, targets)
+        return DatasetSeq(seq_idx, inputFeatures, targets=targets)

{returnn-1.20250226.132109 → returnn-1.20250227.110407}/returnn/frontend/_numpy_backend.py

@@ -159,6 +159,35 @@ class NumpyBackend(Backend[numpy.ndarray]):
             res = numpy.array(res)
         return res
 
+    @staticmethod
+    def where(
+        cond: Tensor,
+        true_: Union[Tensor, rf.RawTensorTypes],
+        false_: Union[Tensor, rf.RawTensorTypes],
+        *,
+        allow_broadcast_all_sources: bool = False,
+    ) -> Tensor:
+        """where"""
+        if isinstance(true_, Tensor):
+            dtype = true_.dtype
+        elif isinstance(false_, Tensor):
+            dtype = false_.dtype
+        else:
+            dtype = None
+        true_ = rf.convert_to_tensor(true_, _backend=NumpyBackend, dtype=dtype)
+        false_ = rf.convert_to_tensor(false_, _backend=NumpyBackend, dtype=dtype)
+        out = Tensor.get_common_data(
+            [true_, false_, cond], allow_broadcast_all_sources=allow_broadcast_all_sources, name="where"
+        )
+        out.dtype = true_.dtype
+        out.sparse_dim = true_.sparse_dim or false_.sparse_dim
+        out.feature_dim = true_.feature_dim or false_.feature_dim
+        cond_bc_raw = cond.copy_compatible_to_dims_raw(out.dims)
+        true_bc_raw = true_.copy_compatible_to_dims_raw(out.dims)
+        false_bc_raw = false_.copy_compatible_to_dims_raw(out.dims)
+        out.raw_tensor = numpy.where(cond_bc_raw, true_bc_raw, false_bc_raw)
+        return out
+
     @staticmethod
     def range_over_dim(dim: Dim, *, dtype: Optional[str] = None, device: Optional[str] = None) -> Tensor[numpy.ndarray]:
         """

{returnn-1.20250226.132109 → returnn-1.20250227.110407}/returnn/tensor/_dim_extra.py

@@ -422,14 +422,10 @@ class _DimMixin:
         :param func: operates inplace
         """
         dyn_size_ext = self.dyn_size_ext.copy() if self.dyn_size_ext is not None else None
-        dyn_size_ext_max = self._dyn_size_max_value if self._dyn_size_max_value is not None else None
         self.reset_raw(only_self=True)
         if dyn_size_ext is not None:
             func(dyn_size_ext)
-        if dyn_size_ext_max is not None:
-            func(dyn_size_ext_max)
         self.dyn_size_ext = dyn_size_ext
-        self._dyn_size_max_value = dyn_size_ext_max
 
     def _can_use_in_ctx(self, ctx):
         """

{returnn-1.20250226.132109 → returnn-1.20250227.110407}/returnn/tensor/_tensor_extra.py

@@ -2966,7 +2966,7 @@ class _TensorMixin(_TensorMixinBase):
         mask = None
         for axis in axes:
             mask_ = self._dims[axis].get_mask(dim_order=self.dims, device=self.device)
-            mask = rf.logical_and(mask, mask_) if mask is not None else mask_
+            mask = rf.combine_bc(mask, "logical_and", mask_) if mask is not None else mask_
         assert isinstance(mask, _t.Tensor)
         res = rf.where(mask, self, mask_value)
         if use_padding_info:

{returnn-1.20250226.132109 → returnn-1.20250227.110407}/returnn/tensor/tensor_dict.py

@@ -9,7 +9,7 @@ We also might have model_outputs in the user config.
 """
 
 from __future__ import annotations
-from typing import Optional, Union, Any, Type, Dict, Sequence
+from typing import Optional, Union, Any, Type, Dict, Sequence, List
 from .tensor import Tensor
 from .dim import Dim
 
@@ -160,6 +160,20 @@ class TensorDict:
                         assert dim.size == raw_tensor_dict[key_]
                 visited_dims.add(dim)
 
+    def all_dims(self) -> List[Dim]:
+        """
+        :return: list of dims
+        """
+        visited_dims = set()
+        out = []
+        for key, value in self.data.items():
+            for dim in value.dims:
+                if dim in visited_dims:
+                    continue
+                out.append(dim)
+                visited_dims.add(dim)
+        return out
+
 
 def _convert_to_tensor(opts: _TensorT, *, name: Optional[str] = None) -> Tensor:
     """

{returnn-1.20250226.132109 → returnn-1.20250227.110407}/returnn/torch/data/pipeline.py

@@ -21,6 +21,7 @@ other PyTorch datasets more directly, including also HuggingFace datasets.
 
 from __future__ import annotations
 import bisect
+import itertools
 from typing import Optional, Any, Sequence, Tuple, Union, List, Dict, Callable
 import sys
 from copy import deepcopy
@@ -65,6 +66,9 @@ def collate_batch(batch: List[Dict[str, numpy.ndarray]]) -> Dict[str, Union[torc
         if key in ("num_seqs", "epoch"):
             res[key] = batch[0][key]  # it should always be the same
             continue
+        elif key == "complete_frac":
+            res[key] = max(sample[key] for sample in batch)
+            continue
         ls = [create_tensor(sample[key]) for sample in batch]
         if not ls:
             raise ValueError("batch is empty?")
@@ -122,7 +126,7 @@ class ChunkingIterDataPipe(torch.utils.data.IterDataPipe):
 
             if not chunking_data_keys:
                 chunking_data_keys = list(data_dict.keys())  # use all if not configured separately
-                chunking_data_key_black_list = ["seq_tag", "seq_idx", "num_seqs", "epoch"]
+                chunking_data_key_black_list = ["seq_tag", "seq_idx", "num_seqs", "epoch", "complete_frac"]
                 for key in chunking_data_key_black_list:
                     if key in chunking_data_keys:
                         chunking_data_keys.remove(key)
@@ -269,8 +273,15 @@ class BatchingIterDataPipe(torch.utils.data.IterDataPipe):
         epoch = int(data_dict["epoch"])
         seq_idx = int(data_dict["seq_idx"])
         num_seqs = int(data_dict["num_seqs"])  # >=1 if known, otherwise -1
-        epoch_continuous = (epoch - 1 + (seq_idx + 1) / num_seqs) if num_seqs > 0 else None
-        return {"epoch": epoch, "seq_idx": seq_idx, "epoch_continuous": epoch_continuous, **get_fwd_compat_kwargs()}
+        complete_frac = float(data_dict["complete_frac"])  # >= 0 if known, otherwise -1
+        epoch_continuous = (
+            epoch - 1 + complete_frac
+            if complete_frac >= 0.0
+            else (epoch - 1 + (seq_idx + 1) / num_seqs)
+            if num_seqs > 0
+            else None
+        )
+        return {"epoch": epoch, "epoch_continuous": epoch_continuous, "seq_idx": seq_idx, **get_fwd_compat_kwargs()}
 
     def __iter__(self):
         """
@@ -455,6 +466,123 @@ class LenFilterDataPipe(torch.utils.data.IterDataPipe):
         raise Exception(f"{self.__class__.__name__}.__getitem__ not supported")
 
 
+class ShufflingDataPipe(torch.utils.data.IterDataPipe):
+    """
+    Data pipe that is similar to ``torch.utils.data.datapipes.iter.Shuffler``,
+    but it will keep certain data keys of the batches in order while shuffling the rest.
+
+    Used for e.g. ``complete_frac`` and ``seq_idx``.
+    """
+
+    def __init__(
+        self,
+        dataset: torch.utils.data.IterableDataset,
+        *,
+        buffer_size: int,
+        monotonic_data_keys: Sequence[str],
+        seed: Optional[int] = None,
+    ):
+        """
+        :param dataset: batches dataset to shuffle
+        :param buffer_size: buffer size for shuffling
+        :param monotonic_data_keys: data keys that will be excluded from shuffling/keep their order
+        :param seed: random seed
+        """
+        super().__init__()
+
+        self._dataset = dataset
+        self._buffer: List[List[Dict[str, Any]]] = []
+        self._next_buffer: List[List[Dict[str, Any]]] = []
+        assert buffer_size > 0
+        self._buffer_size = buffer_size
+        self._monotonic_data_keys = monotonic_data_keys
+        self._rng = numpy.random.RandomState()
+        self._seed = seed
+
+    def __iter__(self):
+        # The implementation is very similar to the PostprocessingDataset's combinator LaplaceOrdering.
+
+        data_iter = iter(self._dataset)
+
+        self._buffer.extend(itertools.islice(data_iter, self._buffer_size))
+        has_ended = False
+
+        while True:
+            # Make sure to not reorder the monotonic values from self._monotonic_data_keys.
+            # These can contain things like complete_frac, which should be kept in order.
+            ordered_data = {
+                key: [data_dict[key] for batch in self._buffer for data_dict in batch]
+                for key in self._monotonic_data_keys
+            }
+            self._rng.shuffle(self._buffer)
+            for key in self._monotonic_data_keys:
+                data_dicts = [data_dict for batch in self._buffer for data_dict in batch]
+                assert len(data_dicts) == len(ordered_data[key])
+                for ordered_value, data_dict in zip(ordered_data[key], data_dicts):
+                    data_dict[key] = ordered_value
+
+            for item in self._buffer:
+                yield item
+
+                try:
+                    if not has_ended:
+                        self._next_buffer.append(next(data_iter))
+                except StopIteration:
+                    has_ended = True
+
+            if len(self._buffer) < self._buffer_size:
+                assert has_ended and not self._next_buffer
+                break
+
+            self._buffer.clear()
+            self._buffer, self._next_buffer = self._next_buffer, self._buffer
+
+    def set_seed(self, seed: int) -> ShufflingDataPipe:
+        """
+        Sets the seed for the next invocation of ``__iter__``, for compatibility with
+        ``torch.utils.data.graph_settings.apply_random_seed``.
+        """
+        self._seed = seed % (2**32)  # seed must be within [0, 2**32) for seeding RandomState
+        return self
+
+    def reset(self):
+        """resets the internal state of the data pipe"""
+        if self._seed is None:
+            self._seed = int(torch.empty((), dtype=torch.int32).random_().item())
+        self._rng.seed(self._seed)
+        self._seed = None
+
+    def __getstate__(self):
+        state = (
+            self._dataset,
+            self._buffer,
+            self._next_buffer,
+            self._buffer_size,
+            self._monotonic_data_keys,
+            self._rng.get_state(),
+            self._seed,
+        )
+        if torch.utils.data.IterDataPipe.getstate_hook is not None:
+            return torch.utils.data.IterDataPipe.getstate_hook(state)
+        return state
+
+    def __setstate__(self, state):
+        (
+            self._dataset,
+            self._buffer,
+            self._next_buffer,
+            self._buffer_size,
+            self._monotonic_data_keys,
+            rng_state,
+            self._seed,
+        ) = state
+        self._rng = numpy.random.RandomState()
+        self._rng.set_state(rng_state)
+
+    def __getitem__(self, index):
+        raise Exception(f"{self.__class__.__name__}.__getitem__ not supported")
+
+
 def create_data_loader_from_batches(
     batches_dataset: torch.utils.data.Dataset, loader_opts: Optional[Dict[str, Any]] = None
 ) -> torch.utils.data.DataLoader: