apply function with compile data

thunder/__init__.py

@@ -824,8 +824,8 @@ def fn_(*args, **kwargs) -> Any:
         cache_entry, inps, pro_to_epi = get_computation_and_inputs(*args, **kwargs)
 
         check_storage_aliases(cache_entry, inps)
-
-        result = cache_entry.computation_fn(*inps)
+        with compile_data_and_stats(cd, cs):
+            result = cache_entry.computation_fn(*inps)
         result = maybe_connect_to_autograd(cache_entry, result)
         result = maybe_call_epilogue(cache_entry, result, pro_to_epi)
 

thunder/core/prims.py

@@ -4030,8 +4030,6 @@ def embedding_backward_meta(grad, indices, num_weights, padding_idx, scale_grad_
 def copy__meta(
     copy_from: TensorProxy,
     copy_to: TensorProxy,
-    *,
-    is_grad_enabled: bool = False,
 ):
     utils.check_type(copy_from, TensorProxy)
     utils.check_type(copy_to, TensorProxy)

thunder/core/transforms.py

@@ -1588,7 +1588,7 @@ def zeros_like(x):
     prims.PrimIDs.LOG2: lambda x: (prims.log2(x), (x,)),
     prims.PrimIDs.ZETA: lambda x, y: (prims.zeta(x, y), (x, y)),
     prims.PrimIDs.FMOD: lambda x, y: (prims.fmod(x, y), (x, y)),
-    prims.PrimIDs.COPY_: lambda x, y, is_grad_enabled: (prims.copy_(x, y, is_grad_enabled=is_grad_enabled), tuple()),
+    prims.PrimIDs.COPY_: lambda x, y: (prims.copy_(x, y), tuple()),
     prims.PrimIDs.CLONE: lambda x: (prims.clone(x), tuple()),
 }
 

thunder/executors/nvfuserex_impl.py

@@ -2054,7 +2054,6 @@ def var_mean(
 def _copy__check(
     copy_from: TensorProxy,
     copy_to: TensorProxy,
-    is_grad_enabled: bool,
 ) -> bool:
     return are_supported_tensors(copy_from, copy_to)
 
@@ -2065,7 +2064,6 @@ def copy_(
     *,
     fd: FusionDefinition,
     lc_to_nv_map: dict,
-    is_grad_enabled: bool,
 ) -> Any:
     nvcopy_from = getnv(copy_from, fd, lc_to_nv_map)
     nvcopy_to = getnv(copy_to, fd, lc_to_nv_map)

thunder/executors/torchex.py

@@ -2181,8 +2181,9 @@ def is_float_type(self, input):
     einops._backends._type2backend[TensorProxy] = EinopsThunderBackend()
 
 
-def _copy__impl(copy_from, copy_to, *, is_grad_enabled):
-    if is_grad_enabled and copy_to.is_leaf and copy_to.requires_grad:
+def _copy__impl(copy_from, copy_to):
+    cd = get_compile_data()
+    if cd.is_grad_enabled and copy_to.is_leaf and copy_to.requires_grad:
         raise RuntimeError("a leaf Variable that requires grad is being used in an in-place operation.")
     copy_to.copy_(copy_from)
     return copy_to

thunder/tests/test_core.py

@@ -350,7 +350,8 @@ def step(self):
 
     optimizer = Optimizer([a, b])
     cstep = executor.make_callable(optimizer.step)
-    cstep()
+    with torch.no_grad():
+        cstep()
 
     expected_a = ref_a - 0.1 * a.grad
     assert_close(a, expected_a)

thunder/tests/test_inplace_functionalization.py

@@ -478,11 +478,11 @@ def f(xs, ys, z):
 def test_inplace_to_tensors_with_grad(executor, device, _):
     @torch.no_grad
     def add_grad(x, y):
-        x.add_(x.grad)
+        return x.add_(x.grad)
 
     @torch.no_grad
     def add_y(x, y):
-        x.add_(y, alpha=0.1)
+        return x.add_(y, alpha=0.1)
 
     for fn in (add_grad, add_y):
         jitted_f = executor.make_callable(fn)
@@ -494,7 +494,8 @@ def add_y(x, y):
         x_ref.grad = x.grad.clone().detach()
         y_ref = y.clone().detach()
 
-        res = jitted_f(x, y)
+        with torch.no_grad():
+            res = jitted_f(x, y)
         res_ref = fn(x_ref, y_ref)
 
         torch.testing.assert_close(x, x_ref)

thunder/torch/__init__.py

@@ -1960,8 +1960,7 @@ def copysign_(a, b, /):
 
 @torchsymbol(torch.Tensor.copy_, is_method=True)  # , tags=(prims.OpTags.IN_PLACE,))
 def copy_(a, b, /):
-    cd = get_compile_data()
-    return prims.copy_(b, a, is_grad_enabled=cd.is_grad_enabled if cd is not None else False)
+    return prims.copy_(b, a)
 
 
 # TODO Implement div