Fixed closures

Author: DoeringChristian

drjit/__init__.py

@@ -2681,15 +2681,15 @@ def inner(closure, *args, **kwargs):
 
         class FrozenFunction:
             def __init__(self, f) -> None:
-                closure = inspect.getclosurevars(f)
-                self.closure = (closure.nonlocals, closure.globals)
+                self.f = f
                 self.frozen = detail.FrozenFunction(
                     inner, limit, warn_after, backend, auto_opaque
                 )
 
             def __call__(self, *args, **kwargs):
                 _state = state_fn(*args, **kwargs) if state_fn is not None else None
-                return self.frozen([self.closure, _state], *args, **kwargs)
+                closure = inspect.getclosurevars(f)
+                return self.frozen([closure.nonlocals, closure.globals, _state], *args, **kwargs)
 
             @property
             def n_recordings(self):
@@ -2725,7 +2725,7 @@ def __get__(self, obj, type=None):
                 if obj is None:
                     return self
                 else:
-                    return FrozenMethod(self.frozen, self.closure, obj)
+                    return FrozenMethod(self.f, self.frozen, obj)
 
         class FrozenMethod(FrozenFunction):
             """
@@ -2738,14 +2738,15 @@ class FrozenMethod(FrozenFunction):
             The ``__call__`` method of the ``FrozenMethod`` then supplies the object
             in addition to the arguments to the internal function.
             """
-            def __init__(self, frozen, closure, obj) -> None:
+            def __init__(self, f, frozen, obj) -> None:
+                self.f = f
                 self.obj = obj
                 self.frozen = frozen
-                self.closure = closure
 
             def __call__(self, *args, **kwargs):
                 _state = state_fn(self.obj, *args, **kwargs) if state_fn is not None else None
-                return self.frozen([self.closure, _state], self.obj, *args, **kwargs)
+                closure = inspect.getclosurevars(self.f)
+                return self.frozen([closure.nonlocals, closure.globals, _state], self.obj, *args, **kwargs)
 
         return functools.wraps(f)(FrozenFunction(f))
 

tests/test_freeze.py

@@ -3085,7 +3085,7 @@ def test80_tensor_mean(t, auto_opaque):
     def func(x):
         return dr.mean(x)
 
-    frozen = dr.freeze(func)
+    frozen = dr.freeze(func, auto_opaque=auto_opaque)
 
     for i in range(3):
         shape = ((i + 3), 10)
@@ -3098,3 +3098,75 @@ def func(x):
 
     assert frozen.n_recordings == 1
 
+
+@pytest.test_arrays("float32, jit, shape=(*)")
+@pytest.mark.parametrize("auto_opaque", [False, True])
+def test81_changing_closures(t, auto_opaque):
+
+    y = 1
+
+    def func(x):
+        return x + y
+
+    frozen  = dr.freeze(func, auto_opaque = auto_opaque)
+
+    for i in range(3):
+        x = dr.arange(t, i + 3)
+
+        res = frozen(x)
+        ref = func(x)
+
+        assert dr.allclose(res, ref)
+
+    assert frozen.n_recordings == 1
+
+    for i in range(3):
+        x = dr.arange(t, i + 3)
+
+        y += 1
+
+        res = frozen(x)
+        ref = func(x)
+
+        assert dr.allclose(res, ref)
+
+    assert frozen.n_recordings == 4
+
+
+@pytest.test_arrays("float32, jit, shape=(*)")
+@pytest.mark.parametrize("auto_opaque", [False, True])
+def test82_changing_closures_methods(t, auto_opaque):
+
+    y = 1
+
+    class Test:
+        def func(self, x):
+            return x + y
+
+        @dr.freeze
+        def frozen(self, x):
+            return x + y
+
+    test = Test()
+
+    for i in range(3):
+        x = dr.arange(t, i + 3)
+
+        res = test.frozen(x)
+        ref = test.func(x)
+
+        assert dr.allclose(res, ref)
+
+    assert test.frozen.n_recordings == 1
+
+    for i in range(3):
+        x = dr.arange(t, i + 3)
+
+        y += 1
+
+        res = test.frozen(x)
+        ref = test.func(x)
+
+        assert dr.allclose(res, ref)
+
+    assert test.frozen.n_recordings == 4