returnn 1.20240712.11947__tar.gz → 1.20240712.164051__tar.gz

{returnn-1.20240712.11947 → returnn-1.20240712.164051}/PKG-INFO

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

returnn-1.20240712.164051/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240712.164051'
+long_version = '1.20240712.164051+git.d5b954b'

{returnn-1.20240712.11947 → returnn-1.20240712.164051}/returnn/frontend/conv.py

@@ -4,6 +4,7 @@ Convolution, transposed convolution, pooling
 
 from __future__ import annotations
 from typing import Optional, Sequence, Tuple, Union
+from returnn.util.basic import next_type_attrib_in_mro_chain
 from returnn.tensor import Tensor, Dim
 import returnn.frontend as rf
 
@@ -98,7 +99,7 @@ class _ConvOrTransposedConv(rf.Module):
         out_spatial_dim: Optional[Dim] = None,
     ) -> Tuple[Tensor, Dim]:
         assert self.nd == 1
-        out, (out_spatial_dim,) = self.__class__.__base__.__call__(
+        out, (out_spatial_dim,) = next_type_attrib_in_mro_chain(self.__class__, "__call__", self.__class__._call_nd1)(
             self,
             source,
             in_spatial_dims=[in_spatial_dim],

{returnn-1.20240712.11947 → returnn-1.20240712.164051}/returnn/torch/util/gradient_checkpoint.py

@@ -530,8 +530,14 @@ def _unregister_custom_saved_tensors_hooks():
         assert thread in _custom_saved_tensors_hooks_registered_threads
         assert (
             not _custom_saved_tensors_hooks_tls_ctx.stack
-            and not _custom_saved_tensors_hooks_tls_ctx.callbacks
+            and (_custom_saved_tensors_hooks_tls_ctx.in_callback or not _custom_saved_tensors_hooks_tls_ctx.callbacks)
             and not _custom_saved_tensors_hooks_tls_ctx.queued_exits
+        ), (
+            f"_unregister_custom_saved_tensors_hooks:"
+            f" stack {_custom_saved_tensors_hooks_tls_ctx.stack},"
+            f" in_callback {_custom_saved_tensors_hooks_tls_ctx.in_callback},"
+            f" callbacks {_custom_saved_tensors_hooks_tls_ctx.callbacks},"
+            f" queued_exits {_custom_saved_tensors_hooks_tls_ctx.queued_exits}"
         )
         _custom_saved_tensors_hooks_tls_ctx.active = False
         _custom_saved_tensors_hooks_registered_threads.remove(thread)
@@ -579,6 +585,7 @@ def _register_custom_saved_tensors_hooks_thread_local_callback(cb: Callable[[],
     Register some thread-local callback function which is called on saved_tensors_hooks __enter__ and __exit__.
     If it returns True, it is kept alive, otherwise removed.
     """
+    assert not _custom_saved_tensors_hooks_tls_ctx.in_callback
     _custom_saved_tensors_hooks_tls_ctx.callbacks.append(cb)
 
 
@@ -590,5 +597,7 @@ def _custom_saved_tensors_hooks_call_callbacks():
         _custom_saved_tensors_hooks_tls_ctx.callbacks = [
             cb for cb in _custom_saved_tensors_hooks_tls_ctx.callbacks if cb()
         ]
+        if not _custom_saved_tensors_hooks_tls_ctx.active:  # was cleaned up by some of the callbacks
+            assert not _custom_saved_tensors_hooks_tls_ctx.callbacks and not _custom_saved_tensors_hooks_tls_ctx.stack
     finally:
         _custom_saved_tensors_hooks_tls_ctx.in_callback = False

{returnn-1.20240712.11947 → returnn-1.20240712.164051}/returnn/util/basic.py

@@ -2396,6 +2396,31 @@ def help_on_type_error_wrong_args(cls, kwargs):
         print("Args mismatch? Missing are %r, unknowns are %r. Kwargs %r." % (mandatory_args, unknown_args, kwargs))
 
 
+def type_attrib_mro_chain(cls: type, attr_name: str) -> list:
+    """
+    :return: list of all attributes with the given name in the MRO chain
+    """
+    attribs = []
+    for cls_ in cls.mro():
+        if attr_name in cls_.__dict__:
+            attribs.append(cls_.__dict__[attr_name])
+    return attribs
+
+
+def next_type_attrib_in_mro_chain(cls: type, attr_name: str, attr):
+    """
+    :param cls:
+    :param attr_name:
+    :param attr: must be in the attrib MRO chain
+    :return: next attribute in the MRO chain
+    """
+    attribs = type_attrib_mro_chain(cls, attr_name)
+    assert attr in attribs
+    idx = attribs.index(attr)
+    assert idx + 1 < len(attribs)
+    return attribs[idx + 1]
+
+
 def custom_exec(source: str, source_filename: str, user_ns: Dict[str, Any], user_global_ns: Dict[str, Any]):
     """
     :param source:

{returnn-1.20240712.11947 → returnn-1.20240712.164051}/returnn.egg-info/PKG-INFO

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

{returnn-1.20240712.11947 → returnn-1.20240712.164051}/tests/test_Util.py

@@ -376,6 +376,23 @@ def test_get_func_kwargs():
     assert_equal(list(getargspec(dummy_func).args), ["net", "var", "update_ops"])
 
 
+def test_next_type_attrib_in_mro_chain():
+    class Base:
+        def method(self):
+            return id(self)
+
+    class Foo(Base):
+        def method(self):
+            return 2
+
+    class Bar(Foo):
+        pass
+
+    assert_equal(type_attrib_mro_chain(Foo, "method"), [Foo.method, Base.method])
+    assert_equal(type_attrib_mro_chain(Bar, "method"), [Foo.method, Base.method])
+    assert_equal(next_type_attrib_in_mro_chain(Bar, "method", Foo.method), Base.method)
+
+
 def test_simple_obj_repr():
     class X:
         def __init__(self, a, b=13):

{returnn-1.20240712.11947 → returnn-1.20240712.164051}/tests/test_rf_base.py

@@ -517,3 +517,9 @@ def test_weight_noise():
                 continue
             print("***", mod, param_name, param)
             rf.weight_noise(mod, param_name, std=0.0025)
+
+    # https://github.com/rwth-i6/returnn/issues/1580
+    conv = rf.Conv1d(in_dim, out_dim, 3, padding="same")
+    rf.weight_noise(conv, "filter", std=0.1)
+    time_dim = Dim(11, name="time")
+    conv(rf.random_normal([time_dim, in_dim]), in_spatial_dim=time_dim)

{returnn-1.20240712.11947 → returnn-1.20240712.164051}/tests/test_torch_util.py

@@ -151,6 +151,51 @@ def test_gradient_checkpoint_scope():
         torch.testing.assert_allclose(param_post_state[k], param_post_state_[k])
 
 
+def test_gradient_checkpoint_scope_twice():
+    # https://github.com/rwth-i6/returnn/issues/1579
+    shape = (101, 103)
+
+    class _Model(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.var = torch.nn.Parameter(torch.randn(shape))
+            self.input_var = torch.nn.Parameter(torch.randn(shape))
+            self.opt = torch.optim.SGD(self.parameters(), lr=0.1)  # not common to have this here but ok for the test
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            return self.get_var() * x
+
+        def get_var(self) -> torch.Tensor:
+            with gradient_checkpoint_scope():
+                return self.var + torch.randn(shape)
+
+        def get_input(self) -> torch.Tensor:
+            x = self.input_var
+            with gradient_checkpoint_scope():
+                return x + torch.randn(shape)
+
+        def demo_run(self):
+            self.opt.zero_grad()
+            y = self(self.get_input())
+            loss = y.sum()  # dummy loss
+            del y  # not needed anymore
+            loss.backward()
+            del loss  # not needed anymore
+            self.opt.step()
+
+    orig_gradient_checkpoint_scope_tensor_del_hook = gradient_checkpoint_scope._tensor_del_hook
+    try:
+        # Overwrite this here to trigger the case where the tensor del hook will not do the cleanup.
+        gradient_checkpoint_scope._tensor_del_hook = lambda self: None
+
+        model = _Model()
+        model.demo_run()
+        model.demo_run()
+
+    finally:
+        gradient_checkpoint_scope._tensor_del_hook = orig_gradient_checkpoint_scope_tensor_del_hook
+
+
 def _report_profile(prof: torch.profiler.profiler, check_events=(), *, _size_threshold=100):
     # Note: I tried prof.events(), prof.profiler.kineto_results.events(), prof._memory_profile().timeline,
     # but they all are not really giving me the information I want.

returnn-1.20240712.11947/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240712.011947'
-long_version = '1.20240712.011947+git.47b8a89'