returnn 1.20250508.93313__py3-none-any.whl → 1.20250513.145447__py3-none-any.whl

returnn/PKG-INFO

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

returnn/_setup_info_generated.py

@@ -1,2 +1,2 @@
-version = '1.20250508.093313'
-long_version = '1.20250508.093313+git.4f05ac7'
+version = '1.20250513.145447'
+long_version = '1.20250513.145447+git.9cdc2a4'

returnn/datasets/basic.py

@@ -20,7 +20,7 @@ import math
 import numpy
 import functools
 import typing
-from typing import TYPE_CHECKING, Optional, Any, Union, Type, Dict, Sequence, List, Callable
+from typing import TYPE_CHECKING, Optional, Any, Set, Tuple, Union, Type, Dict, Sequence, List, Callable
 
 from returnn.log import log
 from returnn.engine.batch import Batch, BatchSetGenerator
@@ -141,12 +141,10 @@ class Dataset:
         :param int _shard_index: local shard index, when sharding is enabled
         """
         self.name = name or ("dataset_id%s" % id(self))
-        self.lock = None  # type: Optional[RLock]  # Used when manipulating our data potentially from multiple threads.
-        self.rnd_seq_drop = None  # type: typing.Optional[Random]
+        self.lock: Optional[RLock] = None  # Used when manipulating our data potentially from multiple threads.
+        self.rnd_seq_drop: Optional[Random] = None
         self.num_inputs = 0  # usually not used, but num_outputs instead, which is more generic
-        self.num_outputs = (
-            None
-        )  # type: typing.Optional[typing.Dict[str,typing.Tuple[int,int]]]  # tuple is num-classes, len(shape).  # nopep8
+        self.num_outputs: Optional[Dict[str, Tuple[int, int]]] = None  # tuple is num-classes, len(shape).
         self.window = window
         self.seq_ordering = seq_ordering  # "default", "sorted" or "random". See self.get_seq_order_for_epoch().
         self.fixed_random_seed = fixed_random_seed
@@ -159,10 +157,10 @@ class Dataset:
         self._seq_order_seq_lens_file = seq_order_seq_lens_file
         self._seq_order_seq_lens_by_idx = None
         # There is probably no use case for combining the two, so avoid potential misconfiguration.
-        assert (
-            self.partition_epoch == 1 or self.repeat_epoch == 1
-        ), "Combining partition_epoch and repeat_epoch is prohibited."
-        self.labels = {}  # type: typing.Dict[str,typing.List[str]]
+        assert self.partition_epoch == 1 or self.repeat_epoch == 1, (
+            "Combining partition_epoch and repeat_epoch is prohibited."
+        )
+        self.labels: Dict[str, List[str]] = {}
         self.weights = {}
         self._num_timesteps = 0
         self._num_seqs = 0
@@ -213,8 +211,8 @@ class Dataset:
             getattr(self, "epoch", "<unknown>"),
         )
 
-    _getnewargs_exclude_attrs = set()  # type: typing.Set[str]
-    _getnewargs_remap = {}  # type: typing.Dict[str,str]
+    _getnewargs_exclude_attrs: Set[str] = set()
+    _getnewargs_remap: Dict[str, str] = {}
 
     @staticmethod
     def _create_from_reduce(cls, kwargs, state) -> Dataset:
@@ -660,12 +658,13 @@ class Dataset:
             )
             old_seq_index = seq_index
             seq_index = [i for i in seq_index if all_seq_tags[i] in self.seq_tags_filter]
-            assert (
-                seq_index
-            ), "%s: empty after applying seq_list_filter_file. Example filter tags: %r, used tags: %r" % (
-                self,
-                sorted(self.seq_tags_filter)[:3],
-                [all_seq_tags[i] for i in old_seq_index[:3]],
+            assert seq_index, (
+                "%s: empty after applying seq_list_filter_file. Example filter tags: %r, used tags: %r"
+                % (
+                    self,
+                    sorted(self.seq_tags_filter)[:3],
+                    [all_seq_tags[i] for i in old_seq_index[:3]],
+                )
             )
         return seq_index
 
@@ -736,9 +735,9 @@ class Dataset:
         """
         self.epoch = epoch
         self.rnd_seq_drop = Random(self._get_random_seed_for_epoch(epoch=epoch))
-        assert (
-            self._num_shards == 1 or self.supports_sharding()
-        ), f"{self}: does not support sharding, but got num_shards == {self._num_shards}"
+        assert self._num_shards == 1 or self.supports_sharding(), (
+            f"{self}: does not support sharding, but got num_shards == {self._num_shards}"
+        )
         return False
 
     def finish_epoch(self, *, free_resources: bool = False):
@@ -970,16 +969,16 @@ class Dataset:
             except Exception:  # also not always available
                 num_seqs = None  # ignore
 
-        if math.isinf(num_seqs):
+        if num_seqs is not None and 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"
+        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

returnn/datasets/cached.py

@@ -46,9 +46,10 @@ class CachedDataset(Dataset):
         self._index_map = range(len(self._seq_index))  # sorted seq idx -> seq_index idx
         self._tag_idx = {}  # type: typing.Dict[str,int]  # map of tag -> real-seq-idx. call _update_tag_idx
         self.targets = {}
-        self.target_keys = (
-            []
-        )  # the keys for which we provide data; we may have labels for additional keys in self.labels
+        # the keys for which we provide data;
+        # we may have labels for additional keys in self.labels
+        self.target_keys = []
+
         self.timestamps = None
 
     def initialize(self):

returnn/datasets/distrib_files.py

@@ -451,8 +451,7 @@ class DistributeFilesDataset(CachedDataset2):
             # We need to decide where to add this file, to the current or the next sub epoch.
             if not files_per_bin[bin_idx] or (
                 # Better to add this file to the current sub epoch?
-                abs((size_taken + size) - avg_size_per_sub_epoch)
-                <= abs(size_taken - avg_size_per_sub_epoch)
+                abs((size_taken + size) - avg_size_per_sub_epoch) <= abs(size_taken - avg_size_per_sub_epoch)
             ):
                 files_per_bin[bin_idx].append(f_tree)
                 size_taken = 0

returnn/datasets/generating.py

@@ -46,12 +46,12 @@ class GeneratingDataset(Dataset):
             output_dim["data"] = (input_dim * self.window, 2)  # not sparse
         self.num_outputs = output_dim
         self.expected_load_seq_start = 0
-        self._seq_order = None  # type: Optional[Sequence[int]]
+        self._seq_order: Optional[Sequence[int]] = None
         self._num_seqs = num_seqs
         self._total_num_seqs = num_seqs
         self.random = numpy.random.RandomState(1)
         self.reached_final_seq = False
-        self.added_data = []  # type: typing.List[DatasetSeq]
+        self.added_data: List[DatasetSeq] = []
         if self.seq_ordering in ("sorted", "sorted_reverse"):
             # For the dev/eval dataset, RETURNN automatically tries to sort them.
             # As this is not supported, just ignore it and reset it to the default order.
@@ -904,22 +904,24 @@ class DummyDatasetMultipleDataKeys(DummyDataset):
             seq_len = {}
             for key in self.data_keys:
                 seq_len[key] = _seq_len
-        assert set(data_keys) == set(
-            seq_len.keys()
-        ), "%s: the keys of seq_len (%s) must match the keys in data_keys=%s." % (
-            self,
-            str(seq_len.keys()),
-            str(data_keys),
+        assert set(data_keys) == set(seq_len.keys()), (
+            "%s: the keys of seq_len (%s) must match the keys in data_keys=%s."
+            % (
+                self,
+                str(seq_len.keys()),
+                str(data_keys),
+            )
+        )
+        assert isinstance(output_dim, dict), (
+            "%s: output_dim %r must be a dict containing a definition for each key in data_keys." % (self, output_dim)
         )
-        assert isinstance(
-            output_dim, dict
-        ), "%s: output_dim %r must be a dict containing a definition for each key in data_keys." % (self, output_dim)
-        assert set(data_keys) == set(
-            output_dim.keys()
-        ), "%s: the keys of output_dim (%s) must match the keys in data_keys=%s." % (
-            self,
-            str(output_dim.keys()),
-            str(data_keys),
+        assert set(data_keys) == set(output_dim.keys()), (
+            "%s: the keys of output_dim (%s) must match the keys in data_keys=%s."
+            % (
+                self,
+                str(output_dim.keys()),
+                str(data_keys),
+            )
         )
 
         super(DummyDatasetMultipleDataKeys, self).__init__(
@@ -2134,9 +2136,7 @@ class LibriSpeechCorpus(CachedDataset2):
         import os
         import zipfile
 
-        transs = (
-            {}
-        )  # type: typing.Dict[typing.Tuple[str,int,int,int],str]  # (subdir, speaker-id, chapter-id, seq-id) -> transcription  # nopep8
+        transs: Dict[Tuple[str, int, int, int], str] = {}  # (subdir, speaker-id, chapter-id, seq-id) -> transcription
         if self.use_zip:
             for name, zip_file in self._zip_files.items():
                 assert isinstance(zip_file, zipfile.ZipFile)

returnn/datasets/hdf.py

@@ -37,9 +37,9 @@ class HDFDataset(CachedDataset):
         :param bool use_cache_manager: uses :func:`Util.cf` for files
         """
         super(HDFDataset, self).__init__(**kwargs)
-        assert (
-            self.partition_epoch == 1 or self.cache_byte_size_total_limit == 0
-        ), "To use partition_epoch in HDFDatasets, disable caching by setting cache_byte_size=0"
+        assert self.partition_epoch == 1 or self.cache_byte_size_total_limit == 0, (
+            "To use partition_epoch in HDFDatasets, disable caching by setting cache_byte_size=0"
+        )
         self._use_cache_manager = use_cache_manager
         self.files = []  # type: typing.List[str]  # file names
         self.h5_files = []  # type: typing.List[h5py.File]
@@ -1246,9 +1246,9 @@ class SimpleHDFWriter:
             self._datasets[name].resize(old_shape[0] + raw_data.shape[0], axis=0)
             expected_shape = (raw_data.shape[0],) + old_shape[1:]
         # append raw data to dataset
-        assert (
-            expected_shape == raw_data.shape
-        ), f"{self} insert: shape mismatch: expected {expected_shape}, got {raw_data.shape}"
+        assert expected_shape == raw_data.shape, (
+            f"{self} insert: shape mismatch: expected {expected_shape}, got {raw_data.shape}"
+        )
         self._datasets[name][self._file.attrs["numTimesteps"] :] = raw_data
         self._file.attrs["numTimesteps"] += raw_data.shape[0]
         self._file.attrs["numSeqs"] += 1
@@ -1302,9 +1302,9 @@ class SimpleHDFWriter:
 
         offset = self._extra_num_time_steps[data_key] - raw_data.shape[0]
         expected_shape = (raw_data.shape[0],) + hdf_data.shape[1:]
-        assert (
-            expected_shape == raw_data.shape
-        ), f"{self} insert other {data_key!r}: shape mismatch: expected {expected_shape}, got {raw_data.shape}"
+        assert expected_shape == raw_data.shape, (
+            f"{self} insert other {data_key!r}: shape mismatch: expected {expected_shape}, got {raw_data.shape}"
+        )
         hdf_data[offset:] = raw_data
 
     def insert_batch(self, inputs, seq_len, seq_tag, extra=None):

returnn/datasets/lm.py

@@ -7,7 +7,22 @@ and some related helpers.
 
 from __future__ import annotations
 
-from typing import Optional, Union, Any, Callable, Iterator, List, Tuple, Set, BinaryIO, Dict, cast, Generator
+from typing import (
+    Iterable,
+    Optional,
+    Sequence,
+    Union,
+    Any,
+    Callable,
+    Iterator,
+    List,
+    Tuple,
+    Set,
+    BinaryIO,
+    Dict,
+    cast,
+    Generator,
+)
 import typing
 import os
 from io import IOBase
@@ -1472,8 +1487,8 @@ class TranslationDataset(CachedDataset2):
         }
 
         self._data_keys = self._source_data_keys + self._target_data_keys
-        self._data = {data_key: [] for data_key in self._data_keys}  # type: typing.Dict[str,typing.List[numpy.ndarray]]
-        self._data_len = None  # type: typing.Optional[int]
+        self._data: Dict[str, List[numpy.ndarray]] = {data_key: [] for data_key in self._data_keys}
+        self._data_len: Optional[int] = None
 
         self._vocabs = self._get_vocabs()
         self.num_outputs = {k: [max(self._vocabs[k].values()) + 1, 1] for k in self._vocabs.keys()}  # all sparse
@@ -1489,7 +1504,7 @@ class TranslationDataset(CachedDataset2):
             unknown_label.setdefault(data_key, None)
         self._unknown_label = unknown_label
 
-        self._seq_order = None  # type: typing.Optional[typing.Sequence[int]]  # seq_idx -> line_nr
+        self._seq_order: Optional[Sequence[int]] = None  # seq_idx -> line_nr
         self._tag_prefix = "line-"  # sequence tag is "line-n", where n is the line number
         self._thread = Thread(name="%r reader" % self, target=self._thread_main)
         self._thread.daemon = True
@@ -1878,14 +1893,11 @@ class TranslationFactorsDataset(TranslationDataset):
             assert file_prefix == self.target_file_prefix
             data_keys = self._target_data_keys
 
-        data = [
+        data: List[List[numpy.ndarray]] = [
             self._factored_words_to_numpy(data_keys, s.decode("utf8").strip().split(), self._add_postfix[file_prefix])
             for s in data_strs
-        ]  # type: typing.List[typing.List[numpy.ndarray]] # shape: (len(data_strs), len(data_keys))
-
-        data = zip(
-            *data
-        )  # type: typing.Iterable[typing.Tuple[numpy.ndarray]] # shape: (len(data_keys), len(data_strs))
+        ]  # shape: (len(data_strs), len(data_keys))
+        data: Iterable[Tuple[numpy.ndarray]] = zip(*data)  # shape: (len(data_keys), len(data_strs))
 
         with self._lock:
             for i, data_ in enumerate(data):
@@ -1908,9 +1920,9 @@ class TranslationFactorsDataset(TranslationDataset):
             words_per_factor = [[]] * len(data_keys)
         elif len(data_keys) > 1:
             factored_words = [word.split(self._factor_separator) for word in words]
-            assert all(
-                len(factors) == len(data_keys) for factors in factored_words
-            ), "All words must have all factors. Expected: " + self._factor_separator.join(data_keys)
+            assert all(len(factors) == len(data_keys) for factors in factored_words), (
+                "All words must have all factors. Expected: " + self._factor_separator.join(data_keys)
+            )
             words_per_factor = zip(*factored_words)
             words_per_factor = [list(w) for w in words_per_factor]
         else:

returnn/datasets/meta.py

@@ -247,10 +247,10 @@ class MetaDataset(CachedDataset2):
         self.seq_order_control_dataset = seq_order_control_dataset
 
         # This will only initialize datasets needed for features occuring in data_map
-        self.datasets = {
+        self.datasets: Dict[str, Dataset] = {
             key: init_dataset(datasets[key], extra_kwargs={"name": "%s_%s" % (self.name, key)}, parent_dataset=self)
             for key in self.dataset_keys
-        }  # type: typing.Dict[str,Dataset]
+        }
 
         self._seq_list_file = seq_list_file
         self.seq_list_original = self._load_seq_list(seq_list_file)
@@ -260,8 +260,8 @@ class MetaDataset(CachedDataset2):
 
         self.tag_idx = {tag: idx for (idx, tag) in enumerate(self.seq_list_original[self.default_dataset_key])}
 
-        self._seq_lens = None  # type: typing.Optional[typing.Dict[str,NumbersDict]]
-        self._num_timesteps = None  # type: typing.Optional[NumbersDict]
+        self._seq_lens: Optional[Dict[str, NumbersDict]] = None
+        self._num_timesteps: Optional[NumbersDict] = None
         self._seq_lens_file = seq_lens_file
         if seq_lens_file:
             seq_lens = load_json(filename=seq_lens_file)
@@ -290,7 +290,7 @@ class MetaDataset(CachedDataset2):
         self.num_outputs = self.data_dims
 
         self.orig_seq_order_is_initialized = False
-        self.seq_list_ordered = None  # type: typing.Optional[typing.Dict[str,typing.List[str]]]
+        self.seq_list_ordered: Optional[Dict[str, List[str]]] = None
 
     def _load_seq_list(self, seq_list_file: Optional[Union[str, Dict[str, str]]] = None) -> Dict[str, List[str]]:
         """
@@ -771,7 +771,7 @@ class ConcatDataset(CachedDataset2):
         for ds in self.datasets[1:]:
             assert ds.num_inputs == self.num_inputs
             assert ds.num_outputs == self.num_outputs
-        self.dataset_seq_idx_offsets = None  # type: typing.Optional[typing.List[int]]
+        self.dataset_seq_idx_offsets: Optional[List[int]] = None
 
     def init_seq_order(self, epoch=None, seq_list=None, seq_order=None):
         """
@@ -1017,9 +1017,9 @@ class CombinedDataset(CachedDataset2):
             for (dset_key, dset_data_key), data_key in data_map.items()
         }
 
-        self.dataset_seq_idx_boundaries = None  # type: typing.Optional[typing.List[int]]
-        self.dataset_sorted_seq_idx_list = None  # type: typing.Optional[typing.List[typing.Tuple[int,int]]]
-        self.used_num_seqs_per_subset = None  # type: typing.Optional[typing.List[int]]
+        self.dataset_seq_idx_boundaries: Optional[List[int]] = None
+        self.dataset_sorted_seq_idx_list: Optional[List[Tuple[int, int]]] = None
+        self.used_num_seqs_per_subset: Optional[List[int]] = None
 
     def init_seq_order(self, epoch=None, seq_list=None, seq_order=None):
         """
@@ -1180,9 +1180,9 @@ class CombinedDataset(CachedDataset2):
         :rtype: list[int]
         """
         assert self.partition_epoch in [None, 1], "partition_epoch not supported in combination with sampling_sizes."
-        assert (
-            self._seq_order_seq_lens_file is None
-        ), "seq_order_seq_lens_file not supported in combination with sampling_sizes."
+        assert self._seq_order_seq_lens_file is None, (
+            "seq_order_seq_lens_file not supported in combination with sampling_sizes."
+        )
         assert not self.unique_seq_tags, "unique_seq_tags not supported in combination with sampling_sizes."
         assert self.seq_tags_filter is None, "seq_order_seq_lens_file in combination with sampling_sizes."
 
@@ -1445,7 +1445,7 @@ class ConcatSeqsDataset(CachedDataset2):
         self.repeat_in_between_last_frame_up_to_multiple_of = repeat_in_between_last_frame_up_to_multiple_of or {}
         self.pad_narrow_data_to_multiple_of_target_len = pad_narrow_data_to_multiple_of_target_len or {}
         if epoch_wise_filter is None:
-            self.epoch_wise_filter = None  # type: Optional[EpochWiseFilter]
+            self.epoch_wise_filter: Optional[EpochWiseFilter] = None
         elif isinstance(epoch_wise_filter, dict):
             self.epoch_wise_filter = EpochWiseFilter(epoch_wise_filter)
         else:
@@ -1471,10 +1471,8 @@ class ConcatSeqsDataset(CachedDataset2):
         self.seq_lens = eval(open(seq_len_file).read())
         assert isinstance(self.seq_lens, dict)
         self.full_seq_len_list = self._get_full_seq_lens_list()
-        self.cur_seq_list = None  # type: typing.Optional[typing.List[str]]  # list of seq tags
-        self.cur_sub_seq_idxs = (
-            None
-        )  # type: typing.Optional[typing.List[typing.List[int]]]  # list of list of sub seq idxs
+        self.cur_seq_list: typing.Optional[typing.List[str]] = None  # list of seq tags
+        self.cur_sub_seq_idxs: typing.Optional[typing.List[typing.List[int]]] = None  # list of list of sub seq idxs
 
     def _get_full_seq_lens_list(self):
         """
@@ -1564,20 +1562,22 @@ class ConcatSeqsDataset(CachedDataset2):
         if seq_idx == 0:  # some extra check, but enough to do for first seq only
             sub_dataset_keys = self.dataset.get_data_keys()
             for key in self.remove_in_between_postfix:
-                assert (
-                    key in sub_dataset_keys
-                ), "%s: remove_in_between_postfix key %r not in sub dataset data-keys %r" % (
-                    self,
-                    key,
-                    sub_dataset_keys,
+                assert key in sub_dataset_keys, (
+                    "%s: remove_in_between_postfix key %r not in sub dataset data-keys %r"
+                    % (
+                        self,
+                        key,
+                        sub_dataset_keys,
+                    )
                 )
             for key in self.repeat_in_between_last_frame_up_to_multiple_of:
-                assert (
-                    key in sub_dataset_keys
-                ), "%s: repeat_in_between_last_frame_up_to_multiple_of key %r not in sub dataset data-keys %r" % (
-                    self,
-                    key,
-                    sub_dataset_keys,
+                assert key in sub_dataset_keys, (
+                    "%s: repeat_in_between_last_frame_up_to_multiple_of key %r not in sub dataset data-keys %r"
+                    % (
+                        self,
+                        key,
+                        sub_dataset_keys,
+                    )
                 )
             for key in self.pad_narrow_data_to_multiple_of_target_len:
                 assert key in sub_dataset_keys, (
@@ -1587,15 +1587,16 @@ class ConcatSeqsDataset(CachedDataset2):
         for sub_seq_idx, sub_seq_tag in zip(sub_seq_idxs, sub_seq_tags):
             self.dataset.load_seqs(sub_seq_idx, sub_seq_idx + 1)
             sub_dataset_tag = self.dataset.get_tag(sub_seq_idx)
-            assert (
-                sub_dataset_tag == sub_seq_tag
-            ), "%s: expected tag %r for sub seq idx %i but got %r, part of seq %i %r" % (
-                self,
-                sub_seq_tag,
-                sub_seq_idx,
-                sub_dataset_tag,
-                seq_idx,
-                seq_tag,
+            assert sub_dataset_tag == sub_seq_tag, (
+                "%s: expected tag %r for sub seq idx %i but got %r, part of seq %i %r"
+                % (
+                    self,
+                    sub_seq_tag,
+                    sub_seq_idx,
+                    sub_dataset_tag,
+                    seq_idx,
+                    seq_tag,
+                )
             )
             for key in self.get_data_keys():
                 data = self.dataset.get_data(sub_seq_idx, key)

returnn/datasets/normalization_data.py

@@ -169,7 +169,7 @@ class NormalizationData:
         sumErr = np.sum(np.abs(newSum - oldSum - intermediateSum))
         if sumErr > NormalizationData.SUMMATION_PRECISION:
             raise FloatingPointError(
-                "sums have very different orders of magnitude." " summation error = {}".format(sumErr)
+                "sums have very different orders of magnitude. summation error = {}".format(sumErr)
             )
         return newSum
 

returnn/datasets/postprocessing.py

@@ -308,19 +308,26 @@ class PostprocessingDataset(CachedDataset2):
                     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 (
-                        in_t.ndim == out_t.batch_ndim
-                        and in_t.dtype == out_t.dtype
-                        and all(d.dimension in (d_, None) for (d, d_) in zip(in_t.dims, out_t.shape))
+                    assert in_t.ndim == out_t.batch_ndim, (
+                        f"Dim number mismatch for {data_key}: {in_t.ndim} != {out_t.batch_ndim}. "
+                        "Postprocessing data tensors must not have a batch dimension."
                     )
+                    assert in_t.dtype == out_t.dtype, (
+                        f"dtype mismatch for {data_key}: '{in_t.dtype}' != '{out_t.dtype}'"
+                    )
+                    for i, (in_dim, out_shape) in enumerate(zip(in_t.dims, out_t.shape)):
+                        assert in_dim.dimension is None or in_dim.dimension == out_shape, (
+                            f"Dim {i} mismatch on {data_key}: "
+                            f"{in_dim.dimension} must either be `None` or equal {out_shape}"
+                        )
                 yield t_dict
 
         data_iter = self._iterate_dataset()
         if self._map_seq_stream is not None:
             data_iter = self._map_seq_stream(data_iter, epoch=self.epoch, rng=self._rng, **util.get_fwd_compat_kwargs())
-            assert isinstance(
-                data_iter, Iterator
-            ), f"map_seq_stream must produce an {Iterator.__name__}, but produced {type(data_iter).__name__}"
+            assert isinstance(data_iter, Iterator), (
+                f"map_seq_stream must produce an {Iterator.__name__}, but produced {type(data_iter).__name__}"
+            )
         return _validate_tensor_dict_iter(data_iter)
 
     def _iterate_dataset(self) -> Iterator[TensorDict]:
@@ -349,9 +356,9 @@ class PostprocessingDataset(CachedDataset2):
                 tensor_dict = self._map_seq(
                     tensor_dict, epoch=self.epoch, seq_idx=seq_index, rng=self._rng, **util.get_fwd_compat_kwargs()
                 )
-                assert isinstance(
-                    tensor_dict, TensorDict
-                ), f"map_seq must produce a {TensorDict.__name__}, but produced {type(tensor_dict).__name__}"
+                assert isinstance(tensor_dict, TensorDict), (
+                    f"map_seq must produce a {TensorDict.__name__}, but produced {type(tensor_dict).__name__}"
+                )
 
                 # 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.
@@ -367,9 +374,9 @@ class PostprocessingDataset(CachedDataset2):
                 if self._seq_list_for_validation is not None:
                     seq_tag = self._seq_list_for_validation[seq_index]
                     tag_of_seq = tensor_dict.data["seq_tag"].raw_tensor.item()
-                    assert (
-                        tag_of_seq == seq_tag
-                    ), f"seq tag mismath: {tag_of_seq} != {seq_tag} for seq index {seq_index} when seq list is given"
+                    assert tag_of_seq == seq_tag, (
+                        f"seq tag mismath: {tag_of_seq} != {seq_tag} for seq index {seq_index} when seq list is given"
+                    )
 
             yield tensor_dict
             seq_index += 1

returnn/datasets/sprint.py

@@ -393,13 +393,14 @@ class SprintDatasetBase(Dataset):
             targets = {"classes": targets}
         if "classes" in targets:
             # 'classes' is always the alignment
-            assert targets["classes"].shape == (
-                reduce_num_frames,
-            ), "Number of targets %s does not match number of features %s (reduce factor %d)" % (
-                # is in format (time,)
-                targets["classes"].shape,
-                (num_frames,),
-                self.reduce_target_factor,
+            assert targets["classes"].shape == (reduce_num_frames,), (
+                "Number of targets %s does not match number of features %s (reduce factor %d)"
+                % (
+                    # is in format (time,)
+                    targets["classes"].shape,
+                    (num_frames,),
+                    self.reduce_target_factor,
+                )
             )
         if "speaker_name" in targets:
             targets["speaker_name"] = targets["speaker_name"].strip()

returnn/datasets/util/strings.py

@@ -2,7 +2,6 @@
 Operations on strings.
 """
 
-
 from __future__ import annotations
 import numpy
 

returnn/datasets/util/vocabulary.py

@@ -185,9 +185,9 @@ class Vocabulary:
                     labels = file_content.splitlines()
                     labels_from_idx = {i: line for (i, line) in enumerate(labels)}
                     labels_to_idx = {line: i for (i, line) in enumerate(labels)}
-            assert isinstance(
-                labels_to_idx, dict
-            ), f"{self}: expected dict, got {type(labels_to_idx).__name__} in {filename}"
+            assert isinstance(labels_to_idx, dict), (
+                f"{self}: expected dict, got {type(labels_to_idx).__name__} in {filename}"
+            )
             if labels_from_idx is None:
                 labels_from_idx = {idx: label for (label, idx) in sorted(labels_to_idx.items())}
             min_label, max_label, num_labels = min(labels_from_idx), max(labels_from_idx), len(labels_from_idx)

returnn/extern/graph_editor/subgraph.py

@@ -12,8 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""SubGraphView: a subgraph view on an existing tf.Graph.
-"""
+"""SubGraphView: a subgraph view on an existing tf.Graph."""
 
 from __future__ import annotations
 

returnn/extern/graph_editor/transform.py

@@ -12,8 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Class to transform an subgraph into another.
-"""
+"""Class to transform an subgraph into another."""
 
 from __future__ import annotations
 

returnn/extern/graph_editor/util.py

@@ -12,8 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Utility functions for the graph_editor.
-"""
+"""Utility functions for the graph_editor."""
 
 from __future__ import annotations
 

returnn/frontend/_backend.py

@@ -1509,9 +1509,10 @@ def get_backend_by_raw_tensor_type(tensor_type: Type[T]) -> Union[Type[Backend[T
         else:
             continue
 
-        assert any(
-            issubclass(base_type, type_) for type_ in tensor_types
-        ), f"tensor type {tensor_type} base_type {base_type} not in {tensor_types}, expected for backend {backend_type}"
+        assert any(issubclass(base_type, type_) for type_ in tensor_types), (
+            f"tensor type {tensor_type} base_type {base_type} not in {tensor_types}, "
+            f"expected for backend {backend_type}"
+        )
         for base_type_ in tensor_types:
             register_backend_by_tensor_type(base_type_, backend_type)
         return backend_type