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

returnn/tf/engine.py

@@ -12,7 +12,7 @@ See :ref:`tech_overview` for an overview how it fits all together.
 
 from __future__ import annotations
 
-from typing import Optional
+from typing import Callable, Dict, List, Optional, Union
 import typing
 import os
 import sys
@@ -101,31 +101,29 @@ class Runner:
         self.store_tf_profile = engine.config.bool("store_tf_profile", False)
         self.store_metadata_mod_step = engine.config.int("store_metadata_mod_step", 0)
         self.reset_updater_vars_mod_step = engine.config.int("reset_updater_vars_mod_step", 0)
-        assert not (
-            self.store_tf_profile and self.store_metadata_mod_step
-        ), "Cannot use store_tf_profile and store_metadata_mod_step at the same time"
+        assert not (self.store_tf_profile and self.store_metadata_mod_step), (
+            "Cannot use store_tf_profile and store_metadata_mod_step at the same time"
+        )
         self.finalized = False
         self.cancel_flag = False
         self.run_exception = None
         self.num_steps = None
-        self.device_crash_batch = None  # type: typing.Optional[int]
+        self.device_crash_batch: Optional[int] = None
         self.start_time = None
         self.elapsed = None
-        self.report_prefix = None  # type: typing.Optional[str]
+        self.report_prefix: Optional[str] = None
         self._results_accumulated = NumbersDict()  # entries like "cost:output" or "loss"
         self._inv_norm_accumulated = NumbersDict()  # entries like "output"
         self.num_frames_accumulated = NumbersDict()  # for each data key (eg. "classes"), corresponding number of frames
-        self.results = {}  # type: typing.Dict[str,float]  # entries like "cost:output" or "loss"
-        self.score = {}  # type: typing.Dict[str,float]  # entries like "cost:output"
-        self.error = {}  # type: typing.Dict[str,float]  # entries like "error:output"
-        self.stats = (
-            {}
-        )  # type: typing.Dict[str,typing.Union[float,numpy.ndarray,'Util.Stats']]  # entries like "stats:..."
+        self.results: Dict[str, float] = {}  # entries like "cost:output" or "loss"
+        self.score: Dict[str, float] = {}  # entries like "cost:output"
+        self.error: Dict[str, float] = {}  # entries like "error:output"
+        self.stats: Dict[str, Union[float, numpy.ndarray, "util.Stats"]] = {}  # entries like "stats:..."
         self.extra_fetches = extra_fetches
         if extra_fetches is not None:
             assert extra_fetches_callback
         self.extra_fetches_callback = extra_fetches_callback
-        self._step_start_time = None  # type: typing.Optional[float]
+        self._step_start_time: Optional[float] = None
         self._horovod_last_param_sync_time = time.time()  # we assume it is synced right now
         self._horovod_stopped_runner = False
         self._horovod_finish_all = False
@@ -133,9 +131,7 @@ class Runner:
             self._horovod_finish_all = True
         # With Horovod, during the main session.run, if reduce_type != grad or not training,
         # the following tensors are enough to ensure that we are in sync.
-        self._horovod_collected_reduce_inputs = (
-            {}
-        )  # type: typing.Dict[str,(tf.Tensor,tf.Tensor)]  # name -> (input,output)
+        self._horovod_collected_reduce_inputs: Dict[str, (tf.Tensor, tf.Tensor)] = {}  # name -> (input,output)
 
         from returnn.util.basic import terminal_size
 
@@ -196,9 +192,9 @@ class Runner:
                     d["extra:%s" % k] = v
                     continue
                 assert isinstance(v, Data)
-                d[
-                    "extra:%s" % k
-                ] = v.placeholder  # see _maybe_handle_extra_fetches, it will transform to batch-major there
+                d["extra:%s" % k] = (
+                    v.placeholder
+                )  # see _maybe_handle_extra_fetches, it will transform to batch-major there
                 for i, s in v.size_placeholder.items():
                     d["extra:%s:size_%i" % (k, i)] = s
 
@@ -732,9 +728,9 @@ class Runner:
                             run_options_.MergeFrom(run_options)
                         # We could use tfdbg.add_debug_tensor_watch here.
                         session_run_start_time = time.time()
-                        fetches_results = sess.run(
+                        fetches_results: Dict[str, Union[numpy.ndarray, str]] = sess.run(
                             fetches_dict, feed_dict=feed_dict, options=run_options_, run_metadata=run_metadata
-                        )  # type: typing.Dict[str,typing.Union[numpy.ndarray,str]]
+                        )
                         elapsed_time_tf += time.time() - session_run_start_time
                         writer.add_summary(fetches_results["summary"], step + step_offset)
                         writer.add_run_metadata(run_metadata, "step_{:04d}".format(step + step_offset))
@@ -746,13 +742,13 @@ class Runner:
                         session_run_start_time = time.time()
                         if self.store_tf_profile:
                             with tf.profiler.experimental.Trace(name=report_prefix, step_num=step + step_offset):
-                                fetches_results = sess.run(
+                                fetches_results: Dict[str, Union[numpy.ndarray, str]] = sess.run(
                                     fetches_dict, feed_dict=feed_dict, options=run_options
-                                )  # type: typing.Dict[str,typing.Union[numpy.ndarray,str]]
+                                )
                         else:
-                            fetches_results = sess.run(
+                            fetches_results: Dict[str, Union[numpy.ndarray, str]] = sess.run(
                                 fetches_dict, feed_dict=feed_dict, options=run_options
-                            )  # type: typing.Dict[str,typing.Union[numpy.ndarray,str]]
+                            )
                         elapsed_time_tf += time.time() - session_run_start_time
                         if writer and "summary" in fetches_results:
                             writer.add_summary(fetches_results["summary"], step + step_offset)
@@ -891,27 +887,27 @@ class Engine(EngineBase):
             BackendEngine.select_engine(default_fallback_engine=default_fallback_engine, config=self.config)
         assert BackendEngine.is_tensorflow_selected()
         self.orig_config = {}  # see _maybe_update_config
-        self.custom_get_net_dict = None  # type: typing.Optional[typing.Callable]
+        self.custom_get_net_dict: Optional[Callable] = None
         self._have_rf_get_model_func = False
         self._check_devices()
-        self.tf_session = None  # type: typing.Optional[tf.compat.v1.Session]
-        self.network = None  # type: typing.Optional[TFNetwork]
-        self.updater = None  # type: typing.Optional[Updater]
+        self.tf_session: Optional[tf.compat.v1.Session] = None
+        self.network: Optional[TFNetwork] = None
+        self.updater: Optional[Updater] = None
         self._checked_uninitialized_vars = False
         self._merge_all_summaries = None
-        self.dataset_batches = {}  # type: typing.Dict[str,BatchSetGenerator]
-        self.dataset_provider = None  # type: typing.Optional[DatasetDataProvider]
-        self.train_data = None  # type: typing.Optional[Dataset]
-        self.eval_datasets = {}  # type: typing.Dict[str,Dataset]
-        self.start_epoch = None  # type: typing.Optional[int]
-        self._num_trained_epochs = 0  # type: int  # just a counter
-        self._num_net_reinit = 0  # type: int
+        self.dataset_batches: Dict[str, BatchSetGenerator] = {}
+        self.dataset_provider: Optional[DatasetDataProvider] = None
+        self.train_data: Optional[Dataset] = None
+        self.eval_datasets: Dict[str, Dataset] = {}
+        self.start_epoch: Optional[int] = None
+        self._num_trained_epochs: int = 0  # just a counter
+        self._num_net_reinit: int = 0
         self.use_dynamic_train_flag = False
         self.use_search_flag = self.config.value("task", None) == "search"
         self.use_eval_flag = self.config.value("task", None) != "forward"
-        self._const_cache = {}  # type: typing.Dict[str,tf.Tensor]
-        self.preload_from_files = None  # type: typing.Optional[typing.Dict[str,typing.Dict[str]]]
-        self.max_seqs = None  # type: typing.Optional[int]
+        self._const_cache: Dict[str, tf.Tensor] = {}
+        self.preload_from_files: Optional[Dict[str, Dict[str]]] = None
+        self.max_seqs: Optional[int] = None
 
     def finalize(self, error_occurred=False):
         """
@@ -1140,7 +1136,7 @@ class Engine(EngineBase):
         self.min_seq_length = config.typed_value("min_seq_length", None) or config.float("min_seq_length", 0)
         self.inc_seq_length = config.float("inc_seq_length", 0)
         if not self.max_seq_length:
-            self.max_seq_length = sys.maxsize  # type: typing.Union[int,float,typing.Dict[str,int],NumbersDict]
+            self.max_seq_length: Union[int, float, Dict[str, int], NumbersDict] = sys.maxsize
         if isinstance(self.max_seq_length, dict):
             self.max_seq_length = NumbersDict(self.max_seq_length)
         assert isinstance(self.max_seq_length, (int, float, NumbersDict))
@@ -1630,7 +1626,7 @@ class Engine(EngineBase):
         assert isinstance(self.start_epoch, int)
         epoch = self.start_epoch  # Epochs start at 1.
         while epoch <= final_epoch:
-            self.epoch = epoch  # type: int
+            self.epoch: int = epoch
             if isinstance(self.max_seq_length, int) and self.max_seq_length != sys.maxsize:
                 if int(self.max_seq_length + self.inc_seq_length) != int(self.max_seq_length):
                     print("increasing sequence lengths to", int(self.max_seq_length + self.inc_seq_length), file=log.v3)
@@ -1878,9 +1874,9 @@ class Engine(EngineBase):
         # We update the model params in-place.
         # In training, we don't want that, because it should not use the validation data.
         # We could reset it later when continuing the training, but it's not implemented.
-        assert (
-            self.config.value("task", "train") != "train"
-        ), "task %r should be just 'eval' or so. training will break." % self.config.value("task", None)
+        assert self.config.value("task", "train") != "train", (
+            "task %r should be just 'eval' or so. training will break." % self.config.value("task", None)
+        )
         if not self.updater:
             self.updater = Updater(
                 config=self.config, network=self.network, initial_learning_rate=self.initial_learning_rate
@@ -1928,11 +1924,12 @@ class Engine(EngineBase):
             allowed_outputs = {"seq_tag", "seq_len", "score", "error", "pos_score", "pos_error"}
 
             assert isinstance(output_per_seq_format, (tuple, list)), "provide output_per_seq_format"
-            assert (
-                set(output_per_seq_format) - allowed_outputs == set()
-            ), "Only %r are allowed in function eval_model as output_per_seq_format, but got: %r " % (
-                allowed_outputs,
-                output_per_seq_format,
+            assert set(output_per_seq_format) - allowed_outputs == set(), (
+                "Only %r are allowed in function eval_model as output_per_seq_format, but got: %r "
+                % (
+                    allowed_outputs,
+                    output_per_seq_format,
+                )
             )
 
             # always fetch seq_tag to map loss values to the corresponding line
@@ -1968,12 +1965,10 @@ class Engine(EngineBase):
             if "pos_error" in output_per_seq_format:
                 extra_fetches["pos_error"] = loss_holder.get_error_value_per_pos()
 
-        seq_idx_to_tag = (
-            {}
-        )  # type: typing.Dict[int,str]  # we need this in order to write the results in the correct order later  # nopep8
-        results_per_seq = (
-            {}
-        )  # type: typing.Dict[str,typing.Dict[str,typing.Union[float,str,int]]]  # seq_tag -> dict. Results of fetches will be written in this dict  # nopep8
+        seq_idx_to_tag: Dict[int, str] = {}  # we need this in order to write the results in the correct order later
+        results_per_seq: Dict[
+            str, Dict[str, Union[float, str, int]]
+        ] = {}  # seq_tag -> dict. Results of fetches will be written in this dict
 
         # function to save the return values of each callback to the dict `results_per_seq`
         # noinspection PyShadowingNames
@@ -2012,7 +2007,7 @@ class Engine(EngineBase):
 
         if output_per_seq_file:
             assert len(self.get_eval_datasets()) == 1, (
-                "output per sequence is only supported for one dataset (dev or eval)," "provided datasets are %r"
+                "output per sequence is only supported for one dataset (dev or eval),provided datasets are %r"
             ) % list(self.get_eval_datasets().keys())
             # try to sort dataset to minimize zero-padding
             dataset = list(self.get_eval_datasets().values())[0]
@@ -2453,9 +2448,9 @@ class Engine(EngineBase):
                 )
 
         max_seq_length = self.config.typed_value("max_seq_length", None) or self.config.float("max_seq_length", 0)
-        assert (
-            not max_seq_length
-        ), "Set max_seq_length = 0 for search (i.e. no maximal length). We want to keep all source sentences."
+        assert not max_seq_length, (
+            "Set max_seq_length = 0 for search (i.e. no maximal length). We want to keep all source sentences."
+        )
 
         dataset.init_seq_order(epoch=self.epoch)
         batches = dataset.generate_batches(
@@ -2552,8 +2547,8 @@ class Engine(EngineBase):
                     outputs[output_layer_idx] = bytearray(outputs[output_layer_idx]).decode("utf8")
 
             # Create lists with serialized data. All of length num_output_layers.
-            serialized_outputs = []  # type: typing.List[typing.Optional[typing.Union[str,numpy.ndarray]]]
-            serialized_targets = []  # type: typing.List[typing.Optional[typing.Union[str,numpy.ndarray]]]
+            serialized_outputs: List[Optional[Union[str, numpy.ndarray]]] = []
+            serialized_targets: List[Optional[Union[str, numpy.ndarray]]] = []
             # noinspection PyShadowingNames
             for output_layer_idx in range(num_output_layers):
                 if output_layers[output_layer_idx].output.sparse:
@@ -2572,8 +2567,8 @@ class Engine(EngineBase):
                         ]
                     else:
                         serialized_output = None
-                        assert not output_file, "Unable to serialize sparse output of layer '%s'." % (
-                            output_layer_names[output_layer_idx]
+                        assert not output_file, (
+                            "Unable to serialize sparse output of layer '%s'." % (output_layer_names[output_layer_idx])
                         )
                 else:
                     # Output dense layers as-is
@@ -2594,8 +2589,8 @@ class Engine(EngineBase):
                             ]
                         else:
                             serialized_target = None
-                            assert not output_file, "Unable to serialize sparse target '%s'." % (
-                                target_keys[output_layer_idx]
+                            assert not output_file, (
+                                "Unable to serialize sparse target '%s'." % (target_keys[output_layer_idx])
                             )
                     else:
                         serialized_target = targets[output_layer_idx]

returnn/tf/frontend_layers/_backend.py

@@ -510,9 +510,9 @@ class ReturnnLayersBackend(Backend[Layer]):
                 # We could also maybe move out all the dependencies.
                 # However, it's not clear whether this is always safe.
                 for dep in value.raw_tensor.get_tensor_dependencies():
-                    assert (
-                        dep.parent.can_access_children_from_root
-                    ), f"dep {dep} of moved value {value} is not accessible"
+                    assert dep.parent.can_access_children_from_root, (
+                        f"dep {dep} of moved value {value} is not accessible"
+                    )
             param.raw_tensor.layer_dict["init_by_layer"] = value
         else:
             param.raw_tensor.layer_dict.pop("init_by_layer", None)

returnn/tf/frontend_layers/cond.py

@@ -181,9 +181,9 @@ class Cond(Generic[T]):
         After this, self.result is available.
         """
         assert self._entered, f"{self} you need to be in the context scope"
-        assert (
-            self._entered_state is False
-        ), f"{self} you need to be in the False branch, have assigned :func:`true` before"
+        assert self._entered_state is False, (
+            f"{self} you need to be in the False branch, have assigned :func:`true` before"
+        )
         assert not self._false_value_set
         nest.assert_same_structure(self._true_value, false_value)
         # This needs to match the true() setter logic.
@@ -198,9 +198,9 @@ class Cond(Generic[T]):
                 if false_v is None:  # see above
                     false_v = rf.zeros((), dtype="int32")  # dummy value
                 else:
-                    assert isinstance(
-                        false_v, Tensor
-                    ), f"unexpected {false_value!r}, only expects tensors, got {type(false_v)}"
+                    assert isinstance(false_v, Tensor), (
+                        f"unexpected {false_value!r}, only expects tensors, got {type(false_v)}"
+                    )
                 assert true_v.raw_tensor.parent is self.true_branch_name_ctx
                 name = true_v.raw_tensor.name
                 assert name not in self.false_branch_name_ctx.children

returnn/tf/frontend_layers/debug_eager_mode.py

@@ -2,7 +2,6 @@
 Debug eager mode
 """
 
-
 _debug_eager_mode_enabled = False
 
 

returnn/tf/frontend_layers/layer.py

@@ -1104,13 +1104,13 @@ class _NetDictBuilderCtx:
                 # If dyn_size_ext is not set yet, try to complete it.
                 if dim.dyn_size_ext is None:
                     dim.complete_dyn_size()
-                assert (
-                    dim.dyn_size_ext is not None
-                ), f"{sub_name_ctx}: need {dim} to be defined to be able to know about implicit dims"
+                assert dim.dyn_size_ext is not None, (
+                    f"{sub_name_ctx}: need {dim} to be defined to be able to know about implicit dims"
+                )
             dim_tags.extend(data_template.dim_tags_set_implicit_only_wrapped)
-            assert len(dim_tags) == len(
-                set((d, d.match_priority if isinstance(d, Dim) else 0) for d in dim_tags)
-            ), f"duplicate dims in {sub_name_ctx} {sub_name_ctx.tensor}"
+            assert len(dim_tags) == len(set((d, d.match_priority if isinstance(d, Dim) else 0) for d in dim_tags)), (
+                f"duplicate dims in {sub_name_ctx} {sub_name_ctx.tensor}"
+            )
             if len(dim_tags) == len(set(dim_tags)):  # might not be unique without match_priority
                 # For some layer classes, the out_shape would be redundant.
                 if layer_dict["class"] not in {"constant", "variable", "random", "subnetwork", "transpose"}:
@@ -1135,9 +1135,9 @@ class _NetDictBuilderCtx:
 
                 sub_layer_abs_name_scope = self._expected_layer_abs_name_scope(sub_name_ctx)
                 if sub_name_ctx.layer_dict["class"] == "variable":
-                    assert (
-                        sub_layer_abs_name_scope
-                    ), f"VariableLayer {sub_name_ctx} must have a unique name in {self.root_module}"
+                    assert sub_layer_abs_name_scope, (
+                        f"VariableLayer {sub_name_ctx} must have a unique name in {self.root_module}"
+                    )
                 if sub_layer_abs_name_scope is not None:
                     if (
                         layer_abs_name_scope_default != sub_layer_abs_name_scope
@@ -1153,9 +1153,9 @@ class _NetDictBuilderCtx:
 
             def _map_elem_resolve(obj: Any) -> Any:
                 if isinstance(obj, Tensor):
-                    assert isinstance(
-                        obj.raw_tensor, rfl.Layer
-                    ), f"unexpected tensor {obj} with raw tensor type {type(obj.raw_tensor)}, expected rfl.Layer"
+                    assert isinstance(obj.raw_tensor, rfl.Layer), (
+                        f"unexpected tensor {obj} with raw tensor type {type(obj.raw_tensor)}, expected rfl.Layer"
+                    )
                     obj: Tensor[rfl.Layer]
                     assert obj.raw_tensor.parent or net.name_ctx == obj.raw_tensor
                     return obj.raw_tensor.get_name_in_ctx(ctx=net.name_ctx)

returnn/tf/frontend_layers/loop.py

@@ -415,9 +415,9 @@ class _LoopState:
             tensor.raw_tensor.make_all_sub_networks_and_optimize()
 
             layer_ctx_list = tensor.raw_tensor.get_abs_name_ctx_list()
-            assert (
-                self.loop.name_ctx in layer_ctx_list
-            ), f"Loop state {name_ctx} should get a value inside the loop but got {tensor}"
+            assert self.loop.name_ctx in layer_ctx_list, (
+                f"Loop state {name_ctx} should get a value inside the loop but got {tensor}"
+            )
             # We need some special logic for MaskedComputation but maybe also for others later.
             # This is currently not nice, but I'm not sure about better solutions.
             for i in range(layer_ctx_list.index(self.loop.name_ctx) + 1, len(layer_ctx_list) - 1):

returnn/tf/frontend_layers/make_layer.py

@@ -74,7 +74,6 @@ def make_layer(
             raise TypeError(f"{layer}: unexpected type {type(value)} in layer_dict: {layer_dict}")
 
     try:
-
         if out is not None:
             assert isinstance(out, Tensor)
         elif predefined_out_data is not None:

returnn/tf/layers/base.py

@@ -4,8 +4,9 @@ This module contains the layer base class :class:`LayerBase`.
 
 from __future__ import annotations
 
-from typing import Optional, Dict, List
+from typing import Optional, Dict, List, Union
 import typing
+from typing import TYPE_CHECKING
 import contextlib
 import numpy
 import tensorflow as tf
@@ -17,6 +18,9 @@ from returnn.tf.util.data import Data, FeatureDim, Dim
 from returnn.tf.util.basic import OutputWithActivation, CustomUpdate, reuse_name_scope
 from returnn.log import log
 
+if TYPE_CHECKING:
+    from tensorflow.python.training.saver import BaseSaverBuilder
+
 
 class LayerBase:
     """
@@ -188,7 +192,7 @@ class LayerBase:
         self.name = name
         self.network = network
         self._register_layer()
-        self.kwargs = None  # type: typing.Optional[typing.Dict[str]] # set via self.post_init
+        self.kwargs: Optional[Dict[str]] = None  # set via self.post_init
         self.target = None
         self.targets = None
         if target:
@@ -219,12 +223,12 @@ class LayerBase:
                 "%s: out_dim handling not implemented correctly for this layer" % self
             )
         out_shape  # noqa  # not used here but in fixup_out_data
-        self.output_before_activation = None  # type: typing.Optional[OutputWithActivation]
-        self.output_loss = None  # type: typing.Optional[tf.Tensor]
+        self.output_before_activation: Optional[OutputWithActivation] = None
+        self.output_loss: Optional[tf.Tensor] = None
         if copy_output_loss_from_source_idx is not None:
             self.output_loss = sources[copy_output_loss_from_source_idx].output_loss
-        self.rec_vars_outputs = {}  # type: typing.Dict[str,tf.Tensor]
-        self.search_choices = None  # type: typing.Optional[SearchChoices]
+        self.rec_vars_outputs: Dict[str, tf.Tensor] = {}
+        self.search_choices: Optional[SearchChoices] = None
         self._src_common_search_choices = _src_common_search_choices
         self._initial_output = initial_output
         self.need_last = need_last
@@ -237,14 +241,14 @@ class LayerBase:
             # Note that this check is somewhat incomplete
             # (does not check multiple sources, see _ConcatInputLayer)
             # and there is no guarantee that a specific layer really uses this correctly.
-            assert sources[0].output.have_dim_tag(
-                in_dim, unique=True
-            ), "%s: in_dim %s not found or unique in input %s" % (self, in_dim, sources[0])
+            assert sources[0].output.have_dim_tag(in_dim, unique=True), (
+                "%s: in_dim %s not found or unique in input %s" % (self, in_dim, sources[0])
+            )
         self.have_params = False
-        self.params = {}  # type: typing.Dict[str,tf.Variable]
-        self.saveable_param_replace = (
-            {}
-        )  # type:  typing.Dict[tf.Variable,typing.Union['tensorflow.python.training.saver.BaseSaverBuilder.SaveableObject',None]]  # see get_saveable_params_dict()  # nopep8
+        self.params: Dict[str, tf.Variable] = {}
+        self.saveable_param_replace: Dict[
+            tf.Variable, Union["BaseSaverBuilder.SaveableObject", None]
+        ] = {}  # see get_saveable_params_dict()
         self.reuse_params = reuse_params
         self.name_scope = name_scope
         self.param_device = param_device
@@ -264,7 +268,7 @@ class LayerBase:
         self.control_dependencies_on_output = control_dependencies_on_output
         self.register_as_extern_data = register_as_extern_data
         # Stats will be collected by the engine.
-        self.stats = {}  # type: typing.Dict[str,tf.Tensor]
+        self.stats: Dict[str, tf.Tensor] = {}
         self._set_prev_state(state)
 
     def _set_prev_state(self, state):
@@ -516,9 +520,9 @@ class LayerBase:
                 # Special case: Input feature or sparse dim looks the same, so overtake it.
                 out_dim = sources_data.feature_dim_or_sparse_dim
         if out_dim:
-            assert (
-                out_dim.dimension == output.dim
-            ), f"Layer {name!r} out_dim {out_dim} does not match Data {output} via out_type {out_type}"
+            assert out_dim.dimension == output.dim, (
+                f"Layer {name!r} out_dim {out_dim} does not match Data {output} via out_type {out_type}"
+            )
             if output.sparse:
                 output.sparse_dim = out_dim
             else:
@@ -850,9 +854,9 @@ class LayerBase:
             loss_scale = d.pop("loss_scale", 1.0)
             if loss_scale != 1.0:
                 if "scale" in loss_opts:
-                    assert (
-                        loss_opts["scale"] == loss_scale
-                    ), "do not use loss_scale and loss with 'scale' option together"
+                    assert loss_opts["scale"] == loss_scale, (
+                        "do not use loss_scale and loss with 'scale' option together"
+                    )
                 loss_opts["scale"] = loss_scale
             d["loss"] = cls._make_loss(
                 class_name=d.pop("loss", None), opts=loss_opts, network=network, get_layer=get_layer
@@ -2099,9 +2103,9 @@ class LayerBase:
                 src_output = src.output.copy()
                 if src_output.placeholder is not None:
                     zeroed_src_shape = tf_util.get_shape(src_output.placeholder)
-                    zeroed_src_shape = [
+                    zeroed_src_shape: List[Union[tf.Tensor, int]] = [
                         zeroed_src_shape[i] for i in range(src_output.batch_ndim)
-                    ]  # type: typing.List[typing.Union[tf.Tensor,int]]
+                    ]
                 else:
                     zeroed_src_shape = []
                     for i, d in enumerate(src_output.batch_shape):
@@ -2550,9 +2554,9 @@ class ReuseParams:
         :rtype: tf.Variable|tf.Tensor
         """
         if self.shape is not None:
-            assert tuple(shape) == tuple(
-                d.dimension for d in self.shape
-            ), "%s: unexpected shape %r for param %r, expected %r" % (self, shape, name, self.shape)
+            assert tuple(shape) == tuple(d.dimension for d in self.shape), (
+                "%s: unexpected shape %r for param %r, expected %r" % (self, shape, name, self.shape)
+            )
         abs_scope_prefix = base_layer.get_absolute_name_scope_prefix()
         assert not abs_scope_prefix or abs_scope_prefix.endswith("/")
         assert name.startswith(abs_scope_prefix)
@@ -2609,10 +2613,10 @@ class SearchChoices:
         assert beam_size is not None
         self.owner = owner
         self._done_src_layer = False
-        self._src_layer = None  # type: typing.Optional[LayerBase]
-        self.src_beams = None  # type: typing.Optional[tf.Tensor]  # src beam index, (batch, beam)
+        self._src_layer: Optional[LayerBase] = None
+        self.src_beams: Optional[tf.Tensor] = None  # src beam index, (batch, beam)
         self.beam_size = beam_size
-        self.beam_scores = None  # type: typing.Optional[tf.Tensor]  # (batch, beam)
+        self.beam_scores: Optional[tf.Tensor] = None  # (batch, beam)
         self.is_decided = is_decided
         self.keep_raw = keep_raw
         if not owner.output.beam:
@@ -2872,22 +2876,22 @@ class Loss:
         """
         self.base_network = base_network
         self.use_flatten_frames = use_flatten_frames
-        self.layer = None  # type: typing.Optional[LayerBase]
+        self.layer: Optional[LayerBase] = None
         # All are initialized in self.init().
-        self.output = None  # type: typing.Optional[Data]
-        self.output_with_activation = None  # type: typing.Optional[OutputWithActivation]
-        self.output_seq_lens = None  # type: typing.Optional[tf.Tensor]
-        self.target = None  # type: typing.Optional[Data]
-        self.target_seq_lens = None  # type: typing.Optional[tf.Tensor]
-        self.output_flat = None  # type: typing.Optional[tf.Tensor]
-        self.output_before_softmax_flat = None  # type: typing.Optional[tf.Tensor]
+        self.output: Optional[Data] = None
+        self.output_with_activation: Optional[OutputWithActivation] = None
+        self.output_seq_lens: Optional[tf.Tensor] = None
+        self.target: Optional[Data] = None
+        self.target_seq_lens: Optional[tf.Tensor] = None
+        self.output_flat: Optional[tf.Tensor] = None
+        self.output_before_softmax_flat: Optional[tf.Tensor] = None
         if _check_output_before_softmax is not None:
             self._check_output_before_softmax = _check_output_before_softmax
-        self.target_flat = None  # type: typing.Optional[tf.Tensor]
+        self.target_flat: Optional[tf.Tensor] = None
         # Maybe make configurable. For now, same as in our Theano behavior.
         # The loss_norm_factor is used by Runner._normalize_loss both for normalization per epoch and per batch.
         # It is e.g. set to 1/sum(target_seq_len), and logic of accumulation is handled in the Runner.
-        self.loss_norm_factor = None  # type: typing.Optional[tf.Tensor]
+        self.loss_norm_factor: Optional[tf.Tensor] = None
         self.use_normalized_loss = use_normalized_loss  # for the optimizer, per batch
         self.custom_norm_factor = custom_norm_factor
         self.custom_inv_norm_factor = custom_inv_norm_factor
@@ -3132,18 +3136,21 @@ class Loss:
                 self.output,
                 self.target,
             )
-        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.get_auto_output_layer_dim(self.target.feature_dim_or_sparse_dim)
-        assert (
-            expected_output_dim.dimension == self.output.dim
-        ), "Expected output dim is %r but the output has dim %r. " % (
-            expected_output_dim,
-            self.output.feature_dim_or_sparse_dim,
-        ) + "Target: %s, output: %s" % (
-            self.target,
-            self.output,
+        assert expected_output_dim.dimension == self.output.dim, (
+            "Expected output dim is %r but the output has dim %r. "
+            % (
+                expected_output_dim,
+                self.output.feature_dim_or_sparse_dim,
+            )
+            + "Target: %s, output: %s"
+            % (
+                self.target,
+                self.output,
+            )
         )
         if self.base_network.get_config().bool("debug_runtime_sanity_checks", False):
             with tf.name_scope("Loss_debug_runtime_sanity_checks"):