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

returnn/tf/layers/basic.py

@@ -4,7 +4,7 @@ Many canonical basic layers.
 
 from __future__ import annotations
 
-from typing import Optional, Union, Sequence, List, Tuple, Dict
+from typing import Callable, Optional, Union, Sequence, List, Tuple, Dict
 import typing
 import tensorflow as tf
 import contextlib
@@ -126,7 +126,7 @@ def concat_sources(src_layers, out_dim=None, allow_broadcast_all_sources=NotSpec
         data.placeholder = tf.concat(
             axis=data.feature_dim_axis, values=[layer_data.placeholder for layer_data in layers_data]
         )
-        axes_split_info = [None] * data.batch_ndim  # type: typing.List[typing.Optional[typing.List[int]]]
+        axes_split_info: List[Optional[List[int]]] = [None] * data.batch_ndim
         axes_split_info[data.feature_dim_axis] = [layer_data.dim for layer_data in layers_data]
         tf_util.set_param_axes_split_info(data.placeholder, axes_split_info)
         # Note: We will loose this info for any further op (e.g. dropout, activation, etc). Should be better...
@@ -294,7 +294,7 @@ class _ConcatInputLayer(LayerBase):
         elif mask == "dropout":
             assert dropout > 0
         self.dropout = dropout
-        self.input_data = None  # type: typing.Optional[Data]
+        self.input_data: Optional[Data] = None
         if self.sources:
             self.input_data = concat_sources_with_opt_dropout(
                 self.sources,
@@ -509,9 +509,7 @@ class ConcatLayer(LayerBase):
         assert sources
         sources, axes = zip(*sources)  # unzip
         axes_int = [layer.output.get_axis_from_description(axis) for (layer, axis) in zip(sources, axes)]
-        concat_dim_tags = [
-            layer.output.dim_tags[axis] for (layer, axis) in zip(sources, axes_int)
-        ]  # type: typing.List[Dim]
+        concat_dim_tags: List[Dim] = [layer.output.dim_tags[axis] for (layer, axis) in zip(sources, axes_int)]
         if any(tag.dimension is None for tag in concat_dim_tags):
             dimension = None
         else:
@@ -707,8 +705,8 @@ class SelectSearchSourcesLayer(InternalLayer):
         self.output = src.output.copy_as_batch_major()
         self.rec_vars_outputs = src.rec_vars_outputs.copy()
         src_search_choices = src.get_search_choices()
-        self.transform_func = None  # type: typing.Optional[typing.Callable[[tf.Tensor],tf.Tensor]]
-        self.search_choices_seq = None  # type: typing.Optional[typing.List[SearchChoices]]
+        self.transform_func: Optional[Callable[[tf.Tensor], tf.Tensor]] = None
+        self.search_choices_seq: Optional[List[SearchChoices]] = None
         if not search_choices:
             assert not src_search_choices
             assert not self.output.beam
@@ -726,13 +724,7 @@ class SelectSearchSourcesLayer(InternalLayer):
             assert src_search_choices in search_choices_seq, self.network.debug_search_choices(
                 self.search_choices_layer
             ) or (
-                (
-                    "%s: No common search base:\n"
-                    "from layer %s\n"
-                    "search choices %s,\n"
-                    "to layer %s\n"
-                    "search choices\n%s."
-                )
+                "%s: No common search base:\nfrom layer %s\nsearch choices %s,\nto layer %s\nsearch choices\n%s."
                 % (self, src, src_search_choices, self.search_choices_layer, pformat(search_choices_seq))
             )
             search_choices_seq = search_choices_seq[: search_choices_seq.index(src_search_choices)]
@@ -4436,12 +4428,13 @@ class MergeDimsLayer(_ConcatInputLayer):
         :rtype: list[int]
         """
         if keep_order:
-            assert isinstance(axes, (tuple, list, typing.Sequence)) and not isinstance(
-                axes, str
-            ), "%s: axes %r must be a list or tuple, to have a well defined order in input %s" % (
-                name,
-                axes,
-                input_data,
+            assert isinstance(axes, (tuple, list, typing.Sequence)) and not isinstance(axes, str), (
+                "%s: axes %r must be a list or tuple, to have a well defined order in input %s"
+                % (
+                    name,
+                    axes,
+                    input_data,
+                )
             )
             axes_ = []
             for axis in axes:
@@ -5562,11 +5555,12 @@ class RepeatLayer(_ConcatInputLayer):
             repetitions_data = repetitions_data.copy_add_dim_by_tag(axis_dim_tag, unbroadcast=True)
         repetitions_axis = repetitions_data.get_axis_from_description(axis, allow_int=False)
         assert repetitions_data.ndim == 1, "Repetitions %r must only have at most one non-batch axis" % repetitions
-        assert (
-            repetitions_data.batch_shape[repetitions_axis] == self.input_data.batch_shape[input_axis]
-        ), "Axis mismatch between input (%i) and repetitions (%i)" % (
-            self.input_data.batch_shape[input_axis],
-            repetitions_data.batch_shape[repetitions_axis],
+        assert repetitions_data.batch_shape[repetitions_axis] == self.input_data.batch_shape[input_axis], (
+            "Axis mismatch between input (%i) and repetitions (%i)"
+            % (
+                self.input_data.batch_shape[input_axis],
+                repetitions_data.batch_shape[repetitions_axis],
+            )
         )
 
         assert self.output.have_batch_axis() == (
@@ -6267,9 +6261,9 @@ class ConvLayer(_ConcatInputLayer):
         from returnn.util import BehaviorVersion
 
         padding = padding.upper() if isinstance(padding, str) else padding
-        assert padding in ["SAME", "VALID", "SAME_STATIC"] or isinstance(
-            padding, (int, tuple, list)
-        ), f"{self}: got unsupported padding {padding}"
+        assert padding in ["SAME", "VALID", "SAME_STATIC"] or isinstance(padding, (int, tuple, list)), (
+            f"{self}: got unsupported padding {padding}"
+        )
         assert "out_type" not in kwargs, "don't set out_type explicitly for this layer"
         assert len(filter_size) in (1, 2, 3), "only 1D conv, 2D conv or 3D conv supported"
         super(ConvLayer, self).__init__(in_dim=in_dim, out_dim=out_dim, **kwargs)
@@ -6285,9 +6279,9 @@ class ConvLayer(_ConcatInputLayer):
         assert len(dilation_rate) == len(filter_size)
         assert not self.input_data.sparse
         assert self.input_data.have_batch_axis()
-        assert (
-            self.input_data.have_feature_axis()
-        ), "this should be our single input feature dim now. otherwise use input_add_feature_dim"
+        assert self.input_data.have_feature_axis(), (
+            "this should be our single input feature dim now. otherwise use input_add_feature_dim"
+        )
         input_data, num_batch_dims = self.transform_input(
             self.input_data,
             network=self.network,
@@ -7117,9 +7111,9 @@ class PoolLayer(_ConcatInputLayer):
         super(PoolLayer, self).__init__(in_dim=in_dim, out_dim=out_dim, **kwargs)
         assert not self.input_data.sparse
         assert self.input_data.have_batch_axis()
-        assert (
-            self.input_data.have_feature_axis()
-        ), "this should be our single input feature dim now. otherwise use input_add_feature_dim"
+        assert self.input_data.have_feature_axis(), (
+            "this should be our single input feature dim now. otherwise use input_add_feature_dim"
+        )
         if in_dim and out_dim:
             assert in_dim == out_dim
         elif in_dim:
@@ -7381,9 +7375,9 @@ class TransposedConvLayer(_ConcatInputLayer):
         out_dim  # noqa  # via get_out_data_from_opts
         assert not self.input_data.sparse
         assert self.input_data.have_batch_axis()
-        assert (
-            self.input_data.have_feature_axis()
-        ), "this should be our single input feature dim now. otherwise use input_add_feature_dim"
+        assert self.input_data.have_feature_axis(), (
+            "this should be our single input feature dim now. otherwise use input_add_feature_dim"
+        )
         input_data, num_batch_dims = ConvLayer.transform_input(
             self.input_data,
             network=self.network,
@@ -7404,14 +7398,15 @@ class TransposedConvLayer(_ConcatInputLayer):
             remove_padding = [remove_padding] * len(spatial_axes)
         if not isinstance(output_padding, (list, tuple)):
             output_padding = [output_padding] * len(spatial_axes)
-        assert (
-            len(spatial_axes) == len(filter_size) == len(strides) == len(remove_padding) == len(output_padding)
-        ), "%s: expected %i-D transposed-conv for input %r but got filter %r and strides %r" % (
-            self,
-            len(spatial_axes),
-            input_data,
-            filter_size,
-            strides,
+        assert len(spatial_axes) == len(filter_size) == len(strides) == len(remove_padding) == len(output_padding), (
+            "%s: expected %i-D transposed-conv for input %r but got filter %r and strides %r"
+            % (
+                self,
+                len(spatial_axes),
+                input_data,
+                filter_size,
+                strides,
+            )
         )
         assert len(spatial_axes) in [1, 2], "%s: %i-D not yet implemented..." % (self, len(spatial_axes))
         x = input_data.placeholder
@@ -8775,9 +8770,9 @@ class DotLayer(LayerBase):
                 red1,
                 red2,
             )
-            assert len(a_reduce_axes) == len(
-                b_reduce_axes
-            ), "%s: sources %r, red1 %r, red2 %r, reduce axes must match in count" % (self, self.sources, red1, red2)
+            assert len(a_reduce_axes) == len(b_reduce_axes), (
+                "%s: sources %r, red1 %r, red2 %r, reduce axes must match in count" % (self, self.sources, red1, red2)
+            )
         if (
             (BehaviorVersion.get() >= 3 and (var1 is NotSpecified or var2 is NotSpecified))
             or var1 == "auto"
@@ -9150,9 +9145,9 @@ class DotLayer(LayerBase):
                 raise Exception(
                     "%s %r: " % (cls.__name__, name) + "%s not found in sources %r" % (red_axis_desc, sources)
                 )
-            assert len(a_reduce_axes) == len(
-                b_reduce_axes
-            ), "%s: sources %r, red1 %r, red2 %r, reduce axes must match in count" % (name, sources, red1, red2)
+            assert len(a_reduce_axes) == len(b_reduce_axes), (
+                "%s: sources %r, red1 %r, red2 %r, reduce axes must match in count" % (name, sources, red1, red2)
+            )
         if (
             (BehaviorVersion.get() >= 3 and (var1 is NotSpecified or var2 is NotSpecified))
             or var1 == "auto"
@@ -10178,9 +10173,9 @@ class CondLayer(LayerBase):
         self.condition_desc = condition
         self.condition_layer = self._make_layer("condition", self.condition_desc)
         self.true_layer_desc = true_layer
-        self.true_layer = None  # type: typing.Optional[LayerBase]
+        self.true_layer: Optional[LayerBase] = None
         self.false_layer_desc = false_layer
-        self.false_layer = None  # type: typing.Optional[LayerBase]
+        self.false_layer: Optional[LayerBase] = None
         assert self.condition_layer.output.batch_ndim == 0 and self.condition_layer.output.dtype == "bool"
         self._extra_out_templates = {k: v[0] for k, v in _extra_out.items()}
         x, extra_out, sizes = tf_util.cond(
@@ -12070,7 +12065,7 @@ class HDFDumpLayer(LayerBase):
             for (key, output) in extra.items()
         }
         extra = {key: output.copy_as_batch_spatial_major() for (key, output) in extra.items()}
-        self.extra = extra  # type: typing.Dict[str,Data]
+        self.extra: Dict[str, Data] = extra
         self.dump_whole_batches = dump_whole_batches
         self.num_seqs_written = 0
         ndim = data.ndim
@@ -12454,9 +12449,9 @@ class BinaryCrossEntropyLoss(Loss):
 
     def _check_init(self):
         assert self.target is not None
-        assert (
-            self.target.batch_ndim == self.output.batch_ndim
-        ), "Number of dimensions mismatch. Target: %s, output: %s" % (self.target, self.output)
+        assert self.target.batch_ndim == self.output.batch_ndim, (
+            "Number of dimensions mismatch. Target: %s, output: %s" % (self.target, self.output)
+        )
 
     def get_value(self):
         """
@@ -13020,7 +13015,7 @@ class ExpectedLoss(Loss):
         self.divide_beam_size = divide_beam_size
         self.subtract_average_loss = subtract_average_loss
         self.loss_correction_grad_only = loss_correction_grad_only
-        self.search_choices = None  # type: typing.Optional[SearchChoices]
+        self.search_choices: Optional[SearchChoices] = None
 
     @classmethod
     def transform_config_dict(cls, d, network, get_layer):
@@ -13120,9 +13115,9 @@ class DeepClusteringLoss(Loss):
         Does some checks on self.target and self.output, e.g. if the dense shapes matches.
         You can overwrite this if those checks don't make sense for your derived loss class.
         """
-        assert (
-            self.target.ndim_dense == self.output.ndim_dense
-        ), "Number of dimensions mismatch. Target: %s, output: %s" % (self.target, self.output)
+        assert self.target.ndim_dense == self.output.ndim_dense, (
+            "Number of dimensions mismatch. Target: %s, output: %s" % (self.target, self.output)
+        )
         expected_output_dim = self._embedding_dimension * (self.target.shape[1] // self._nr_of_sources)
         assert expected_output_dim == self.output.dim, "Expected output dim is %i but the output has dim %r. " % (
             expected_output_dim,
@@ -13822,9 +13817,9 @@ class SamplingBasedLoss(Loss):
                     else:
                         loss_fn = tf.nn.sampled_softmax_loss
 
-                    assert (
-                        self.layer.params["W"].shape[0] == self.target.dim
-                    ), "Expect weight matrix of shape [num_classes, dim]"
+                    assert self.layer.params["W"].shape[0] == self.target.dim, (
+                        "Expect weight matrix of shape [num_classes, dim]"
+                    )
                     out = loss_fn(
                         weights=self.layer.params["W"].read_value(),  # (num_classes,D).
                         biases=self.layer.params["b"].read_value(),  # (num_classes).

returnn/tf/layers/rec.py

@@ -6,6 +6,7 @@ from __future__ import annotations
 
 import contextlib
 import typing
+from typing import Dict, Optional, Tuple, Union
 import tensorflow as tf
 import returnn.tf.compat as tf_compat
 
@@ -1037,7 +1038,8 @@ class RecLayer(_ConcatInputLayer):
                     scope=tf_compat.v1.get_variable_scope(),
                 )
         elif rnn_contrib and isinstance(
-            cell, (rnn_contrib.FusedRNNCell, rnn_contrib.LSTMBlockWrapper)  # noqa # e.g. LSTMBlockFusedCell
+            cell,
+            (rnn_contrib.FusedRNNCell, rnn_contrib.LSTMBlockWrapper),  # noqa # e.g. LSTMBlockFusedCell
         ):
             # Will get (time,batch,ydim).
             assert self._max_seq_len is None
@@ -1280,9 +1282,9 @@ class RecLayer(_ConcatInputLayer):
         :param str|int|None key:
         :rtype: tf.Tensor
         """
-        assert (
-            self._last_hidden_state is not None
-        ), "last-hidden-state not implemented/supported for this layer-type. try another unit. see the code."
+        assert self._last_hidden_state is not None, (
+            "last-hidden-state not implemented/supported for this layer-type. try another unit. see the code."
+        )
         return RnnCellLayer.get_state_by_key(self._last_hidden_state, key=key)
 
     @classmethod
@@ -1431,9 +1433,7 @@ class _SubnetworkRecCell:
             )
         self._last_frames = {}  # type: typing.Dict[str,Data]
         self._initial_outputs = None  # type: typing.Optional[typing.Dict[str,tf.Tensor]]
-        self._initial_extra_outputs = (
-            None
-        )  # type: typing.Optional[typing.Dict[str,typing.Dict[str,typing.Union[tf.Tensor,typing.Tuple[tf.Tensor,...]]]]]  # nopep8
+        self._initial_extra_outputs: Optional[Dict[str, Dict[str, Union[tf.Tensor, Tuple[tf.Tensor, ...]]]]] = None
 
         # input_layers_moved_out, output_layers_moved_out and layers_in_loop include (used) sub-layers as separate
         # entries, this way in- and outputting them to the loop via TensorArrays will be handled just as for normal
@@ -1608,14 +1608,9 @@ class _SubnetworkRecCell:
                 while parent and parent.parent:
                     parent_names.insert(0, parent.parent_name or "?")
                     parent = parent.parent
-                return (
-                    "<RecLayer construct template GetLayer>("
-                    "allow_uninitialized_template %r, "
-                    "parents %r)"
-                    % (
-                        lself.allow_uninitialized_template,
-                        " <- ".join(parent_names) or None,
-                    )
+                return "<RecLayer construct template GetLayer>(allow_uninitialized_template %r, parents %r)" % (
+                    lself.allow_uninitialized_template,
+                    " <- ".join(parent_names) or None,
                 )
 
             def _add_uninitialized_count(self):
@@ -2141,16 +2136,17 @@ class _SubnetworkRecCell:
             layer = self.input_layers_net.layers[layer_name]
             assert isinstance(layer, LayerBase)
             if layer_name not in inputs_moved_out_tas:
-                assert not layer.output.mark_same_time(
-                    self._time_dim_tags
-                ), "%s does not expect to have matching time dim to %s" % (layer, self.parent_rec_layer)
-                assert (
-                    name != "output" and not prev
-                ), "Time dim does not match: RecLayer %s (%r) vs sub layer %s (%r)." % (
-                    self.parent_rec_layer,
-                    self.parent_rec_layer.output.get_time_dim_tag(),
-                    layer,
-                    layer.output.get_time_dim_tag(),
+                assert not layer.output.mark_same_time(self._time_dim_tags), (
+                    "%s does not expect to have matching time dim to %s" % (layer, self.parent_rec_layer)
+                )
+                assert name != "output" and not prev, (
+                    "Time dim does not match: RecLayer %s (%r) vs sub layer %s (%r)."
+                    % (
+                        self.parent_rec_layer,
+                        self.parent_rec_layer.output.get_time_dim_tag(),
+                        layer,
+                        layer.output.get_time_dim_tag(),
+                    )
                 )
                 return layer
             output = layer.output.copy_template_excluding_time_dim().copy_template_set_ctx(self.net.control_flow_ctx)
@@ -2376,9 +2372,9 @@ class _SubnetworkRecCell:
         assert output_template.output.dim == self.parent_rec_layer.output.dim
         assert self.parent_rec_layer.output.time_dim_axis == 0
         assert not output_template.output.has_axis(self.time_dim_tag)
-        assert (
-            output_template.output.batch_shape == self.parent_rec_layer.output.batch_shape[1:]
-        ), "see RecLayer.get_out_data_from_opts()"
+        assert output_template.output.batch_shape == self.parent_rec_layer.output.batch_shape[1:], (
+            "see RecLayer.get_out_data_from_opts()"
+        )
 
     def get_init_loop_vars(self):
         """
@@ -3014,9 +3010,9 @@ class _SubnetworkRecCell:
                 needed_outputs.add("end")
                 assert tf.as_dtype(end_template.output.dtype) is tf.bool
             else:
-                assert (
-                    have_known_seq_len
-                ), "You need to have an 'end' layer in your rec subnet if the generated seq len is unknown."
+                assert have_known_seq_len, (
+                    "You need to have an 'end' layer in your rec subnet if the generated seq len is unknown."
+                )
 
             # noinspection PyProtectedMember
             if self.parent_rec_layer._optimize_move_layers_out:
@@ -3358,11 +3354,12 @@ class _SubnetworkRecCell:
                         from .basic import SelectSearchSourcesLayer
 
                         prev_end_layer = choices.translate_to_this_search_beam(prev_end_layer)
-                        assert isinstance(
-                            prev_end_layer, SelectSearchSourcesLayer
-                        ), "unexpected search choices: cur end %r, prev end %r" % (
-                            choices,
-                            prev_end_layer.get_search_choices(),
+                        assert isinstance(prev_end_layer, SelectSearchSourcesLayer), (
+                            "unexpected search choices: cur end %r, prev end %r"
+                            % (
+                                choices,
+                                prev_end_layer.get_search_choices(),
+                            )
                         )
                         prev_end_flag = prev_end_layer.output.placeholder
                         with tf.name_scope("dyn_seq_len"):
@@ -3475,14 +3472,15 @@ class _SubnetworkRecCell:
                 assert fixed_seq_len is not None
                 seq_len = fixed_seq_len
                 if output_beam:
-                    assert (
-                        not input_beam or input_beam == output_beam
-                    ), "%s: input beam %r, output beam %r, sources %r, target %r" % (
-                        self.parent_rec_layer,
-                        input_beam,
-                        output_beam,
-                        self.parent_rec_layer.sources,
-                        self.parent_rec_layer.target,
+                    assert not input_beam or input_beam == output_beam, (
+                        "%s: input beam %r, output beam %r, sources %r, target %r"
+                        % (
+                            self.parent_rec_layer,
+                            input_beam,
+                            output_beam,
+                            self.parent_rec_layer.sources,
+                            self.parent_rec_layer.target,
+                        )
                     )
                     assert output_template.output.batch.beam == output_beam
                     time_dim_tag = time_dim_tag.get_for_batch_ctx(
@@ -3791,9 +3789,9 @@ class _SubnetworkRecCell:
             if end_layer_choice.name.startswith("prev:"):
                 # Logic from maybe_transform. It would be translated to the current beam.
                 end_layer_choice = self.net.layers[end_layer_choice.name[len("prev:") :]]
-            assert (
-                end_layer_choice in choice_seq_in_frame
-            ), "End layer must not have a beam independent from output layer '{}'.".format(layer_name)
+            assert end_layer_choice in choice_seq_in_frame, (
+                "End layer must not have a beam independent from output layer '{}'.".format(layer_name)
+            )
 
             end_layer_choice_index = choice_seq_in_frame.index(end_layer_choice)
             choices_seq_until_end_layer = choice_seq_in_frame[:end_layer_choice_index]
@@ -5856,12 +5854,13 @@ class RecUnstackLayer(LayerBase):
                 if out_dim.is_dim_known():  # usually the case except at template construction
                     assert out_dim != rec_time_dim  # rec_time_dim is unknown, so it cannot be the same
                 if out_dim != rec_time_dim:
-                    assert (
-                        declare_rec_time
-                    ), "%s %r: must either set known axis on rec %s or enable declare_rec_time" % (
-                        cls.__name__,
-                        name,
-                        rec_time_dim,
+                    assert declare_rec_time, (
+                        "%s %r: must either set known axis on rec %s or enable declare_rec_time"
+                        % (
+                            cls.__name__,
+                            name,
+                            rec_time_dim,
+                        )
                     )
                     rec_time_dim.declare_same_as(out_dim)
             out.mark_same_time(out_dim, must_match=True)
@@ -6132,12 +6131,13 @@ class ChoiceLayer(BaseChoiceLayer):
                 base_beam_in = tf.shape(scores_base)[1]  # 1 in first frame, then beam_in
                 scores_beam_in = tf.shape(scores_in)[0] // net_batch_dim
                 beam_in = self.sources[0].output.beam.beam_size
-                assert (
-                    beam_in == base_search_choices.beam_size
-                ), "%r: source %r beam-size unexpected from base choice %r" % (
-                    self,
-                    self.sources[0],
-                    base_search_choices,
+                assert beam_in == base_search_choices.beam_size, (
+                    "%r: source %r beam-size unexpected from base choice %r"
+                    % (
+                        self,
+                        self.sources[0],
+                        base_search_choices,
+                    )
                 )
                 # About incoming beam size:
                 #   base_beam_in  - 1 in first frame, then beam_in
@@ -7510,9 +7510,9 @@ class GenericAttentionLayer(AttentionBaseLayer):
         base_rem_axes = base.get_axes(exclude_batch=True, exclude_time=True)
         base_rem_axes.remove(base.feature_dim_axis)
         weights_rem_axes = weights.get_axes(exclude_batch=True)
-        assert (
-            weights.time_dim_axis is not None
-        ), f"{exception_prefix}: base {base}, weights {weights}, need time_dim_axis in weights"
+        assert weights.time_dim_axis is not None, (
+            f"{exception_prefix}: base {base}, weights {weights}, need time_dim_axis in weights"
+        )
         weights_axis_to_reduce = cls._weights_time_axis_to_reduce(weights=weights, base=base)
         assert weights.batch_shape[weights_axis_to_reduce] == base.batch_shape[base.time_dim_axis]
         weights_rem_axes.remove(weights_axis_to_reduce)
@@ -9088,13 +9088,13 @@ class MaskedComputationLayer(LayerBase):
         new_size, new_time, idxs = None, None, None
         if mask:
             if self.network.is_inside_rec_layer():
-                assert (
-                    mask.output.shape == () and mask.output.dtype == "bool"
-                ), "%s: invalid mask %s (inside rec loop)" % (self, mask)
+                assert mask.output.shape == () and mask.output.dtype == "bool", (
+                    "%s: invalid mask %s (inside rec loop)" % (self, mask)
+                )
             else:
-                assert (
-                    mask.output.have_time_axis() and mask.output.shape == (None,) and mask.output.dtype == "bool"
-                ), "%s: invalid mask %s (outside rec loop)" % (self, mask)
+                assert mask.output.have_time_axis() and mask.output.shape == (None,) and mask.output.dtype == "bool", (
+                    "%s: invalid mask %s (outside rec loop)" % (self, mask)
+                )
                 assert in_spatial_dim and out_spatial_dim
                 mask_data = mask.output.copy_as_time_major()
                 mask_t = where_bc(mask_data.placeholder, mask_data.get_sequence_mask(), tf.convert_to_tensor(False))
@@ -9785,9 +9785,9 @@ class UnmaskLayer(LayerBase):
                 with same_control_flow_ctx(src_layer.output.placeholder):
                     src = src_layer.output.copy_as_bt_or_tb_major()
                 mask_out = self.mask.output
-                assert (
-                    mask_out.shape == () and mask_out.batch_shape == (None,) and mask_out.dtype == "bool"
-                ), "%s: invalid mask %s (inside rec loop)" % (self, self.mask)
+                assert mask_out.shape == () and mask_out.batch_shape == (None,) and mask_out.dtype == "bool", (
+                    "%s: invalid mask %s (inside rec loop)" % (self, self.mask)
+                )
                 prev_t = self._rec_previous_layer.rec_vars_outputs["t"]  # [B]
                 t = prev_t + tf.cast(mask_out.placeholder, tf.int32)  # [B]
                 self.rec_vars_outputs["t"] = t
@@ -11192,9 +11192,9 @@ class RelativePositionalEncodingLayer(_ConcatInputLayer):
             and is_axis_from_description_recurrent(key_value_spatial_dim, network=self.network, data=self.input_data)
         ):
             length = self.network.get_rec_step_index() + 1
-            assert (
-                key_value_spatial_dim_.dimension is None
-            ), f"{self}: unexpected kv spatial dim {key_value_spatial_dim_}"
+            assert key_value_spatial_dim_.dimension is None, (
+                f"{self}: unexpected kv spatial dim {key_value_spatial_dim_}"
+            )
             assert key_value_spatial_dim_.dyn_size_ext is not None
             # See CumConcatLayer for similar logic
             if key_value_spatial_dim_.dyn_size_ext.placeholder is None:

returnn/tf/native_op.py

@@ -283,9 +283,7 @@ class OpMaker:
           //   otherwise it will trigger an assertion.
           if (IsRefType(context->input_dtype({in_idx})))
             context->forward_ref_input_to_ref_output({in_idx}, {out_idx});
-          """.format(
-                    in_idx=in_idx, out_idx=out_idx
-                )
+          """.format(in_idx=in_idx, out_idx=out_idx)
         code_set_io = ""
         for in_idx, v in enumerate(in_info):
             ndim = len(v["shape"])