returnn 1.20240219.173458__tar.gz → 1.20240306.151239__tar.gz

{returnn-1.20240219.173458/returnn.egg-info → returnn-1.20240306.151239}/PKG-INFO

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

returnn-1.20240306.151239/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240306.151239'
+long_version = '1.20240306.151239+git.9b9b07b'

{returnn-1.20240219.173458 → returnn-1.20240306.151239}/returnn/frontend/run_ctx.py

@@ -214,7 +214,7 @@ class RunCtx:
             custom_inv_norm_factor=custom_inv_norm_factor,
         )
 
-    def mark_as_output(self, tensor: Union[Tensor, Any], name: str, *, dims: Optional[Sequence[int]] = None) -> None:
+    def mark_as_output(self, tensor: Union[Tensor, Any], name: str, *, dims: Optional[Sequence[Dim]] = None) -> None:
         """
         Mark this as an output.
         This has the effect that RETURNN will in any case construct the corresponding layer.

{returnn-1.20240219.173458 → returnn-1.20240306.151239}/returnn/torch/engine.py

@@ -384,9 +384,11 @@ class Engine(EngineBase):
             total_loss = train_ctx.total_loss()
             losses_dict = NumbersDict(
                 {
-                    name: float(loss.get_summed_loss().raw_tensor.detach().cpu().numpy())
-                    if self._device != "meta"
-                    else float("nan")
+                    name: (
+                        float(loss.get_summed_loss().raw_tensor.detach().cpu().numpy())
+                        if self._device != "meta"
+                        else float("nan")
+                    )
                     for name, loss in train_ctx.losses.items()
                 }
             )
@@ -400,9 +402,11 @@ class Engine(EngineBase):
                 )
             cur_count_grad_accum += 1
             perform_update_step = cur_count_grad_accum >= accum_grad_multiple_step
-            with self._ddp_pt_model.no_sync() if (
-                self._ddp_pt_model is not None and not perform_update_step
-            ) else nullcontext():
+            with (
+                self._ddp_pt_model.no_sync()
+                if (self._ddp_pt_model is not None and not perform_update_step)
+                else nullcontext()
+            ):
                 if self._grad_scaler is not None:
                     self._grad_scaler.scale(total_loss.raw_tensor).backward()
                 else:
@@ -522,9 +526,11 @@ class Engine(EngineBase):
 
                     losses_dict = NumbersDict(
                         {
-                            name: float(loss.get_summed_loss().raw_tensor.detach().cpu().numpy())
-                            if self._device != "meta"
-                            else float("nan")
+                            name: (
+                                float(loss.get_summed_loss().raw_tensor.detach().cpu().numpy())
+                                if self._device != "meta"
+                                else float("nan")
+                            )
                             for name, loss in train_ctx.losses.items()
                         }
                     )
@@ -636,9 +642,11 @@ class Engine(EngineBase):
                 expected_outputs=self._forward_step_expected_outputs, step=self.global_train_step, epoch=self.epoch
             )
 
-        with autocast(
-            device_type=self._device.split(":")[0], dtype=self._autocast_dtype
-        ) if self._use_autocast else nullcontext(), rf.set_default_device_ctx(self._device):
+        with (
+            autocast(device_type=self._device.split(":")[0], dtype=self._autocast_dtype)
+            if self._use_autocast
+            else nullcontext()
+        ), rf.set_default_device_ctx(self._device):
             sentinel_kw = {"__fwd_compatible_random_arg_%i" % int(random() * 100): None}
             if train_func:
                 self._train_step_func(model=self._orig_model, extern_data=extern_data, **sentinel_kw)
@@ -761,6 +769,17 @@ class Engine(EngineBase):
         self.epoch = epoch  # in training, this will be reset to start_epoch
         self.global_train_step = step
 
+        load_model_post_hooks = self.config.typed_value("load_model_post_hooks")
+        if load_model_post_hooks:
+            with (
+                autocast(device_type=self._device.split(":")[0], dtype=self._autocast_dtype)
+                if self._use_autocast
+                else nullcontext()
+            ), rf.set_default_device_ctx(self._device):
+                sentinel_kw = {"__fwd_compatible_random_arg_%i" % int(random() * 100): None}
+                for hook in load_model_post_hooks:
+                    hook(model=self._orig_model, **sentinel_kw)
+
     def _create_model(self, *, epoch: int, step: int):
         """
         Set up self._pt_model and self._orig_model

{returnn-1.20240219.173458 → returnn-1.20240306.151239}/returnn/torch/frontend/_backend.py

@@ -393,12 +393,13 @@ class TorchBackend(Backend[torch.Tensor]):
         if allow_broadcast:
             for source, dim in sources[1:]:
                 assert dim in source.dims
+                if set(source.dims) - {dim} != set(other_dims):
+                    need_broadcast = True
                 for dim_ in source.dims:
                     if dim_ == dim:
                         continue
                     if dim_ not in other_dims:
                         other_dims.append(dim_)
-                        need_broadcast = True
         sources_raw = []
         if allow_broadcast and need_broadcast:
             for source, dim in sources:

{returnn-1.20240219.173458 → returnn-1.20240306.151239}/returnn/util/basic.py

@@ -1086,9 +1086,9 @@ class DictAsObj:
     See also :class:`ObjAsDict`.
     """
 
-    def __init__(self, dikt):
+    def __init__(self, dikt: Dict[str, Any]):
         """
-        :param dict[str] dikt:
+        :param dikt:
         """
         self.__dict__ = dikt
 
@@ -2690,10 +2690,10 @@ class CollectionReadCheckCovered:
     to check whether all the user-specified options are also used (maybe the user made a typo).
     """
 
-    def __init__(self, collection, truth_value=None):
+    def __init__(self, collection: Dict[str, Any], truth_value: Optional[bool] = None):
         """
-        :param dict[str] collection:
-        :param None|bool truth_value: note: check explicitly for self.truth_value, bool(self) is not the same!
+        :param collection:
+        :param truth_value: note: check explicitly for self.truth_value, bool(self) is not the same!
         """
         self.collection = collection
         if truth_value is None:
@@ -2760,6 +2760,7 @@ def which(program):
     :return: full path, e.g. "/usr/bin/python", or None
     :rtype: str|None
     """
+
     # noinspection PyShadowingNames
     def is_exe(path):
         """
@@ -4066,7 +4067,7 @@ def compute_bleu(reference_corpus, translation_corpus, max_order=4, use_bp=True)
     matches_by_order = [0] * max_order
     possible_matches_by_order = [0] * max_order
 
-    for (references, translations) in zip(reference_corpus, translation_corpus):
+    for references, translations in zip(reference_corpus, translation_corpus):
         reference_length += len(references)
         translation_length += len(translations)
         ref_ngram_counts = _get_ngrams(references, max_order)

{returnn-1.20240219.173458 → returnn-1.20240306.151239/returnn.egg-info}/PKG-INFO

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

returnn-1.20240219.173458/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240219.173458'
-long_version = '1.20240219.173458+git.f7a807f'