returnn 1.20240906.140550__tar.gz → 1.20240906.150016__tar.gz

{returnn-1.20240906.140550 → returnn-1.20240906.150016}/PKG-INFO

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

returnn-1.20240906.150016/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240906.150016'
+long_version = '1.20240906.150016+git.7a58317'

{returnn-1.20240906.140550 → returnn-1.20240906.150016}/returnn/tf/network.py

@@ -3,7 +3,7 @@ Defines the :class:`TFNetwork` and :class:`ExternData`.
 """
 
 from __future__ import annotations
-from typing import Optional, List, Tuple
+from typing import Optional, Any, List, Tuple, Dict
 import tensorflow as tf
 import sys
 import re
@@ -650,8 +650,8 @@ class TFNetwork(object):
         self.subnets = {}  # type: typing.Dict[str,Subnetwork]
         self._selected_train_layers = None
         self._construction_stack = _NetworkConstructionStack()
-        self.layers_desc = {}  # type: typing.Dict[str,typing.Dict[str]]
-        self.layers = {}  # type: typing.Dict[str,LayerBase]
+        self.layers_desc: Dict[str, Dict[str, Any]] = {}
+        self.layers: Dict[str, LayerBase] = {}
         self.losses_dict = {}  # type: typing.Dict[str,LossHolder]
         self.total_loss = None  # type: typing.Optional[tf.Tensor]
         self.total_constraints = None  # type: typing.Optional[tf.Tensor]
@@ -663,16 +663,16 @@ class TFNetwork(object):
         self.extra_vars_to_save = []  # type: typing.List[tf.Variable]
         self.recurrent = False
         self._assigner_cache = {}  # type: typing.Dict[tf.Variable,tf_util.VariableAssigner]
-        self.concat_sources_dropout_cache = (
-            {}
-        )  # type: typing.Dict[typing.Tuple[typing.Tuple[LayerBase,...],Dim,float,typing.Optional[typing.Tuple[typing.Optional[int],...]]],Data]  # nopep8
+        self.concat_sources_dropout_cache: Dict[
+            Tuple[Tuple[LayerBase, ...], Dim, float, Optional[Tuple[Optional[int], ...]]], Data
+        ] = {}
         self._merge_all_summaries = None  # type: typing.Optional[tf.Tensor]
         self._graph_reset_callbacks = []  # type: typing.List[typing.Callable]
         self._run_opts = {}  # type: typing.Dict[str, typing.Any]
         self._run_finished_callbacks = []  # type: typing.List[typing.Callable]
         self._map_search_beam_to_search_choices = (
             {}
-        )  # type: typing.Dict[tf_util.SearchBeam,"returnn.tf.layers.base.SearchChoices"]  # nopep8
+        )  # type: typing.Dict[tf_util.SearchBeam,"returnn.tf.layers.base.SearchChoices"]
 
     def __repr__(self):
         s = "TFNetwork %r" % self.name
@@ -786,18 +786,17 @@ class TFNetwork(object):
             return self.extra_parent_net.get_absolute_name_prefix() + my_prefix
         return ""
 
-    def construct_from_dict(self, net_dict, get_layer=None):
+    def construct_from_dict(self, net_dict: Dict[str, Dict[str, Any]], get_layer=None):
         """
-        :param dict[str,dict[str]] net_dict:
+        :param net_dict:
         :param GetLayer|((str)->LayerBase)|None get_layer:
         """
         self.layers_desc.update(net_dict)
 
-        def ignore_layer(name_, layer_desc_):
+        def ignore_layer(name_: str, layer_desc_: Dict[str, Any]) -> bool:
             """
-            :param str name_:
-            :param dict layer_desc_:
-            :rtype: bool
+            :param name_:
+            :param layer_desc_:
             """
             assert isinstance(name_, str)
             if name_.startswith("#"):
@@ -1028,14 +1027,21 @@ class TFNetwork(object):
         """
         return self.get_config().bool("flat_net_construction", False)
 
-    def construct_layer(self, net_dict, name, get_layer=None, add_layer=None, check_existing=True):
+    def construct_layer(
+        self,
+        net_dict: Dict[str, Dict[str, Any]],
+        name: str,
+        get_layer=None,
+        add_layer=None,
+        check_existing: bool = True,
+    ) -> LayerBase:
         """
         This triggers the construction of the layer `name` if it is not constructed yet.
         Every construction trigger corresponds to ``add_layer`` call (which by default does the actual construction).
         This can recursively also get/construct other layers (via ``get_layer``).
 
-        :param dict[str,dict[str]] net_dict:
-        :param str name: layer name
+        :param net_dict:
+        :param name: layer name
         :param GetLayer|((str)->LayerBase)|None get_layer: optional, for source layers, for transform_config_dict.
           By default, this wraps self.construct_layer().
           I.e. the name might be misleading, as this should return an existing layer,
@@ -1050,8 +1056,8 @@ class TFNetwork(object):
             But otherwise, this is not an option!
 
         :param ((str, LayerBase, dict) -> LayerBase) | None add_layer: by default self.add_layer
-        :param bool check_existing: check self.get_layer. (self.layers will be checked in any case)
-        :rtype: LayerBase
+        :param check_existing: check self.get_layer. (self.layers will be checked in any case)
+        :return: layer
         """
         if name in self.layers:
             return self.layers[name]

{returnn-1.20240906.140550 → returnn-1.20240906.150016}/returnn/tf/util/basic.py

@@ -4,7 +4,7 @@ Lots of random utility functions for TensorFlow.
 
 from __future__ import annotations
 
-from typing import Optional, Union, Sequence, List
+from typing import Optional, Union, Any, Sequence, List, Dict
 import typing
 import contextlib
 from dataclasses import dataclass
@@ -1658,14 +1658,19 @@ class GammatoneFilterbankInitializer(init_ops.Initializer):
         return tf.convert_to_tensor(fbank, dtype=dtype)
 
 
-def get_initializer(s, seed=None, eval_local_ns=None, dtype=tf.float32):
+def get_initializer(
+    s,
+    seed: Optional[Union[int, tf.Tensor]] = None,
+    eval_local_ns: Optional[Dict[str, Any]] = None,
+    dtype: Union[tf.DType, str] = tf.float32,
+):
     """
     :param str|dict[str]|float|numpy.ndarray s: e.g. "glorot_uniform" or "truncated_normal" or "orthogonal",
       or config dict with "class",
       or string to be `eval`ed if it contains "(". constant if a float is given.
-    :param int|tf.Tensor seed: used in case the initializer has no explicit seed specified.
-    :param dict[str]|None eval_local_ns:
-    :param tf.DType|str dtype:
+    :param seed: used in case the initializer has no explicit seed specified.
+    :param eval_local_ns:
+    :param dtype:
     :return: (function (shape) -> tf.Tensor) | tf.Initializer
     :rtype: ((tuple[int]) -> tf.Tensor) | tf.Initializer
     """

{returnn-1.20240906.140550 → returnn-1.20240906.150016}/returnn/tf/util/gradient_checkpoint.py

@@ -73,21 +73,53 @@ def prepare_gradient_checkpointing():
     from tensorflow.python.framework import ops
 
     copied_ops = {}  # type: typing.Dict[tf.Operation, tf.Operation]  # old -> new
+    copy_op_stack_depth = 0
+
+    class _DeepCopyError(Exception):
+        # noinspection PyShadowingNames
+        def __init__(self, op: tf.Operation):
+            super().__init__(f"deep copy err: {op}")
+            self.op = op
 
     # noinspection PyShadowingNames
     def _copy_op(op: tf.Operation) -> tf.Operation:
         if op in copied_ops:
             return copied_ops[op]
 
-        new_inputs = []
-        for x in op.inputs:
-            x = _map_tensor(x)
-            new_inputs.append(x)
+        nonlocal copy_op_stack_depth
+        if copy_op_stack_depth >= 1:
+            # Avoid deep recursions here, as this can get very deep on big graphs.
+            # So do a flat construction in the loop below.
+            raise _DeepCopyError(op)
+
+        try:
+            copy_op_stack_depth += 1
+
+            new_op = None
+            copy_op_queue = [op]
+            while copy_op_queue:
+                op = copy_op_queue[-1]
+
+                try:
+                    new_inputs = []
+                    for x in op.inputs:
+                        x = _map_tensor(x)
+                        new_inputs.append(x)
+
+                except _DeepCopyError as exc:
+                    copy_op_queue.append(exc.op)
+                    continue
+
+                with tf_util.same_control_flow_ctx(op.outputs[0]), tf.name_scope(""):
+                    new_op = tf_util.copy_op(op, inputs=new_inputs, name=op.name)
+                _set_wrapped_grad_func(new_op)
+                copied_ops[op] = new_op
+
+                assert op is copy_op_queue[-1]
+                copy_op_queue.pop(-1)
 
-        with tf_util.same_control_flow_ctx(op.outputs[0]), tf.name_scope(""):
-            new_op = tf_util.copy_op(op, inputs=new_inputs, name=op.name)
-        _set_wrapped_grad_func(new_op)
-        copied_ops[op] = new_op
+        finally:
+            copy_op_stack_depth -= 1
         return new_op
 
     def _map_tensor(x: tf.Tensor) -> tf.Tensor:

{returnn-1.20240906.140550 → returnn-1.20240906.150016}/returnn.egg-info/PKG-INFO

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

returnn-1.20240906.140550/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240906.140550'
-long_version = '1.20240906.140550+git.fd95009'