returnn 1.20240714.161251__tar.gz → 1.20240718.93955__tar.gz

{returnn-1.20240714.161251 → returnn-1.20240718.93955}/PKG-INFO

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

returnn-1.20240718.93955/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240718.093955'
+long_version = '1.20240718.093955+git.cf394a2'

{returnn-1.20240714.161251 → returnn-1.20240718.93955}/returnn/frontend/array_.py

@@ -558,8 +558,12 @@ def sequence_mask(dims: Union[Dim, Sequence[Dim]], *, device: Optional[str] = No
         dims = [dims]
     assert len(dims) > 0
     dyn_dims = [d for d in dims if d.need_masking()]
-    assert len(dyn_dims) == 1  # not implemented otherwise yet...
-    return dyn_dims[0].get_mask(dim_order=dims, device=device)
+    if not dyn_dims:
+        return rf.constant(True, dims=())
+    mask = True
+    for dim in dyn_dims:
+        mask = rf.opt_logical_and(mask, dim.get_mask(dim_order=dims, device=device))
+    return mask
 
 
 def pack_padded(

{returnn-1.20240714.161251 → returnn-1.20240718.93955}/returnn/frontend/math_.py

@@ -36,6 +36,7 @@ __all__ = [
     "logical_or",
     "logical_not",
     "opt_logical_or",
+    "opt_logical_and",
     "maximum",
     "minimum",
     "clip_by_value",
@@ -341,6 +342,24 @@ def opt_logical_or(a: Union[Tensor, bool], b: Union[Tensor, bool]) -> Union[Tens
     return combine(a, "logical_or", b)
 
 
+@overload
+def opt_logical_and(a: bool, b: bool) -> bool:
+    """logical and"""
+
+
+def opt_logical_and(a: Union[Tensor, bool], b: Union[Tensor, bool]) -> Union[Tensor, bool]:
+    """logical and"""
+    if isinstance(a, bool):
+        if not a:
+            return False
+        return b
+    if isinstance(b, bool):
+        if not b:
+            return False
+        return a
+    return combine(a, "logical_and", b)
+
+
 def maximum(a: Tensor, b: Union[Tensor, _RawTensorTypes], *other_tensors) -> Tensor:
     """maximum"""
     if not other_tensors:

{returnn-1.20240714.161251 → returnn-1.20240718.93955}/returnn/torch/frontend/_backend.py

@@ -529,10 +529,10 @@ class TorchBackend(Backend[torch.Tensor]):
         f = getattr(torch, func, None)
         if not f:
             f = getattr(torch.nn.functional, func, None)
-        if not f:
-            f = TorchBackend._ActivationFuncMapping.get(func)
-        if not f:
-            raise ValueError(f"unknown activation function {func!r}")
+            if not f:
+                f = TorchBackend._ActivationFuncMapping.get(func)
+                if not f:
+                    raise ValueError(f"unknown activation function {func!r}")
         return f(raw_tensor)
 
     @staticmethod

{returnn-1.20240714.161251 → returnn-1.20240718.93955}/returnn/torch/util/gradient_checkpoint.py

@@ -206,7 +206,8 @@ class gradient_checkpoint_scope:
             # so we can exit saved_tensors_hooks_scope now.
             # (We might not always catch this properly in the Tensor _DelHook,
             #  e.g. when Tensor.__del__ runs in a different thread.)
-            self.exit_saved_tensors_hooks_scope()
+            if _can_exit_saved_tensors_hooks_inside_hooks():
+                self.exit_saved_tensors_hooks_scope()
             return x
         # _RecordGraph.__torch_dispatch__ should have recorded all newly created tensors.
         x_ = self.record_graph_scope.graph.graph_tensor_from_weak_raw_tensor.get(x, x)
@@ -217,7 +218,8 @@ class gradient_checkpoint_scope:
     @staticmethod
     def _unpack_hook(x: Union[torch.Tensor, _GraphTensor]) -> torch.Tensor:
         if isinstance(x, _GraphTensor):
-            x.op.graph.gradient_checkpoint_scope_backref._maybe_exit_saved_tensors_hooks_scope()
+            if _can_exit_saved_tensors_hooks_inside_hooks():
+                x.op.graph.gradient_checkpoint_scope_backref._maybe_exit_saved_tensors_hooks_scope()
             x.op.graph.maybe_recompute()
             return x.get_recomputed()
         return x
@@ -495,6 +497,18 @@ class _WeakMethod:
         return self.func(obj, *args, **kwargs)
 
 
+def _can_exit_saved_tensors_hooks_inside_hooks() -> bool:
+    """
+    Check whether we can call torch.autograd.graph.saved_tensors_hooks.__exit__
+    inside the pack or unpack hook itself.
+
+    https://github.com/rwth-i6/returnn/issues/1581
+    https://github.com/pytorch/pytorch/issues/130734
+    """
+    # TODO Return True in some later PyTorch version? Check https://github.com/pytorch/pytorch/issues/130734.
+    return False
+
+
 _orig_saved_tensors_hooks_enter = torch.autograd.graph.saved_tensors_hooks.__enter__
 _orig_saved_tensors_hooks_exit = torch.autograd.graph.saved_tensors_hooks.__exit__
 _custom_saved_tensors_hooks_tls_ctx = threading.local()
@@ -559,7 +573,9 @@ def _custom_saved_tensors_hooks_enter(self: torch.autograd.graph.saved_tensors_h
     return _orig_saved_tensors_hooks_enter(self)
 
 
-def _custom_saved_tensors_hooks_exit(self: torch.autograd.graph.saved_tensors_hooks, exc_type, exc_val, exc_tb):
+def _custom_saved_tensors_hooks_exit(
+    self: torch.autograd.graph.saved_tensors_hooks, exc_type=None, exc_val=None, exc_tb=None
+):
     if self not in _custom_saved_tensors_hooks_tls_ctx.stack:
         raise Exception(
             f"saved_tensors_hooks __exit__ mismatch."

{returnn-1.20240714.161251 → returnn-1.20240718.93955}/returnn.egg-info/PKG-INFO

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

{returnn-1.20240714.161251 → returnn-1.20240718.93955}/tests/test_torch_util.py

@@ -12,13 +12,14 @@ import unittest
 import torch
 
 from returnn.util import better_exchook
-from returnn.torch.util.gradient_checkpoint import gradient_checkpoint_scope
 
 
+@unittest.skipIf(torch.__version__ < (2,), "gradient_checkpoint_scope needs PyTorch >= 2.0")
 def test_gradient_checkpoint_scope():
     # https://github.com/rwth-i6/returnn/issues/1552
     from copy import deepcopy
     from torch.profiler import profile, record_function, ProfilerActivity
+    from returnn.torch.util.gradient_checkpoint import gradient_checkpoint_scope
 
     shape = (101, 103)
 
@@ -151,8 +152,12 @@ def test_gradient_checkpoint_scope():
         torch.testing.assert_allclose(param_post_state[k], param_post_state_[k])
 
 
+@unittest.skipIf(torch.__version__ < (2,), "gradient_checkpoint_scope needs PyTorch >= 2.0")
 def test_gradient_checkpoint_scope_twice():
     # https://github.com/rwth-i6/returnn/issues/1579
+
+    from returnn.torch.util.gradient_checkpoint import gradient_checkpoint_scope
+
     shape = (101, 103)
 
     class _Model(torch.nn.Module):
@@ -196,6 +201,50 @@ def test_gradient_checkpoint_scope_twice():
         gradient_checkpoint_scope._tensor_del_hook = orig_gradient_checkpoint_scope_tensor_del_hook
 
 
+@unittest.skipIf(torch.__version__ < (2,), "gradient_checkpoint_scope needs PyTorch >= 2.0")
+def test_saved_tensors_hooks_gc_segfault():
+    # https://github.com/rwth-i6/returnn/issues/1581
+    # https://github.com/pytorch/pytorch/issues/130734
+
+    # noinspection PyProtectedMember
+    from returnn.torch.util.gradient_checkpoint import _can_exit_saved_tensors_hooks_inside_hooks
+
+    if not _can_exit_saved_tensors_hooks_inside_hooks():
+        raise unittest.SkipTest("Not yet fixed.")
+
+    shape = (101, 103)
+    for i in range(10):
+        print("**** iter", i)
+        v = torch.nn.Parameter(torch.randn(shape))
+
+        class _Handler:
+            def __init__(self):
+                self.scope = torch.autograd.graph.saved_tensors_hooks(self._pack_hook, self._unpack_hook)
+                self.scope.__enter__()
+                self.exited = False
+
+            def _pack_hook(self, x):
+                print(f"*** _pack_hook {self}")
+                return x
+
+            def _unpack_hook(self, x):
+                print(f"*** _unpack_hook {self}")
+                if not self.exited:
+                    self.exited = True
+                    print(
+                        f"*** exit {self.scope},"
+                        f" pack_hook {hex(id(self.scope.pack_hook))},"
+                        f" unpack_hook {hex(id(self.scope.unpack_hook))}"
+                    )
+                    self.scope.__exit__()
+                return x
+
+        with torch.autograd.graph.saved_tensors_hooks(lambda x: x, lambda x: x):
+            handler = _Handler()  # keep ref...  # noqa
+            x = v * torch.randn(shape)
+            x.sum().backward()
+
+
 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.20240714.161251/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240714.161251'
-long_version = '1.20240714.161251+git.2e54d6a'