Can I use grad around a numpy code and compile it?

Roy-Kid: <p>Hi,</p>
<p>I have a set of pure numpy codes, and I want to calculate the gradient.<br>
Can I do something like:</p>
<pre><code class="lang-auto">Z = grad(numpy_func)(x)
</code></pre>
<p>Since <code>torch.compile</code> can understand numpy code, and also it can work with <code>grad</code>(<a href="https://pytorch.org/docs/main/torch.compiler_faq.html#does-torch-func-work-with-torch-compile-for-grad-and-vmap-transforms" class="inline-onebox" rel="noopener nofollow ugc">Frequently Asked Questions — PyTorch main documentation</a>). I wonder if I can use grad(and other func) on a numpy function?</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/roy-kid">@Roy-Kid</a>,</p>
<p>In order to compute gradients of a function within the pytorch framework, you need to 1) Use <code>torch</code> functions so that autograd can track operations or 2) Define a custom <code>torch.autograd.Function</code> and manually define the backward pass of your function so autograd can compute gradients.</p>
<p>The docs for <code>torch.autograd.Function</code> objects can be found here: <a href="https://pytorch.org/docs/stable/autograd.html#torch.autograd.Function" class="inline-onebox" rel="noopener nofollow ugc">Automatic differentiation package - torch.autograd — PyTorch 2.3 documentation</a></p>

Roy-Kid: <p>Hi,</p>
<p>I want to use <code>torch.func.grad</code> to calculate the gradient not autograd. The numpy function can be compiled to torch code; I don’t think I need to implement a backward process manually?</p>

AlphaBetaGamma96: <p>I still think you need to define the backward manually, can you take the <code>grad</code> of a compiled function from <code>numpy</code>?</p>
<p>You could always try doing <code>torch.func.grad(compiled_fn)</code> and seeing how <code>torch.func</code> works, but I feel it won’t work as <code>torch.func</code> works on pytorch primitives.</p>

Roy-Kid: <p>I hope so… According to <a href="https://pytorch.org/docs/main/torch.compiler_faq.html#does-torch-func-work-with-torch-compile-for-grad-and-vmap-transforms" class="inline-onebox" rel="noopener nofollow ugc">Frequently Asked Questions — PyTorch main documentation</a>, it can use torch.func.grad and compile together(I test it in 2.4.0-dev).</p>
<p>What if I want to do following:</p>
<pre><code class="lang-auto">import torch

@torch.compile
def f(x):
    return np.sin(x)

def g(x):
    return torch.func.grad(f)(x)

x = torch.randn(2, 3)
g(x)
</code></pre>
<p>The error raises: <code>RuntimeError: Can't call numpy() on Tensor that requires grad. Use tensor.detach().numpy() instead.</code> So I want to ask if it not support yet or I can not do it at all.</p>

AlphaBetaGamma96: <p>Can you share a minimal reproducible example with your function?</p>
<p>Also, from the example you shared it only uses <code>torch</code> functions,</p>
<pre><code class="lang-auto">import torch

@torch.compile
def f(x):
    return torch.sin(x)

def g(x):
    return torch.grad(f)(x)

x = torch.randn(2, 3)
g(x)
</code></pre>
<p>What version of pytorch are you using? The <a href="https://pytorch.org/docs/main/torch.compiler_faq.html#does-numpy-work-with-torch-compile" rel="noopener nofollow ugc">docs</a> state it needs to be 2.1+</p>

Roy-Kid: <p>This one is tested under 2.4.0-dev</p>
<pre><code class="lang-auto">import torch
import numpy as np
from torch import compile
from torch.func import grad

def fn(x):
    return np.power(x, 2).sum()

cfn = compile(fn)
gcfn = grad(compile(fn))
cgfn = compile(grad(fn))
cgcfn = compile(grad(compile(fn)))

arr = torch.tensor([1., 2., 3.])
print(fn(arr))
print(cfn(arr))
# print(gcfn(arr))  error
# print(cgfn(arr))  error
print(cgcfn(arr))  # wrong
</code></pre>

AlphaBetaGamma96: <p>So, perhaps you might have to pass a <code>full_graph=True</code> flag to the compiling statement, in the <a href="https://pytorch.org/docs/main/torch.compiler_faq.html#can-i-execute-numpy-code-on-cuda-and-compute-gradients-via-torch-compile" rel="noopener nofollow ugc">docs</a>, it has this following example,</p>
<pre><code class="lang-auto">@torch.compile(fullgraph=True)
@torch.compiler.wrap_numpy
def numpy_fn(X, Y):
    return np.sum(X[:, :, None] * Y[:, None, :], axis=(-2, -1))

X = torch.randn(1024, 64, device="cuda", requires_grad=True)
Y = torch.randn(1024, 64, device="cuda")
Z = numpy_fn(X, Y)
assert isinstance(Z, torch.Tensor)
Z.backward()
# X.grad now holds the gradient of the computation
print(X.grad)
</code></pre>

Roy-Kid: <p>Although this snippet can not run in 2.4.0-dev(I think it also can not run in a stable version because numpy_fn’s return is not scalar), it looks more reasonable.</p>
<p>I need to use <code>torch.func.grad</code>, not <code>Tensor.grad</code>, because it is more convenient and easy to calculate high-order differences. I think it should work but it does not:</p>
<pre><code class="lang-auto">@torch.compile(fullgraph=True)
@torch.compiler.wrap_numpy
def numpy_fn(X, Y):
    return np.sum(X[:, :, None] * Y[:, None, :], axis=(-2, -1))

X = torch.randn(1024, 64, device="cuda")
Y = torch.randn(1024, 64, device="cuda")
Z = torch.func.grad(numpy_fn)(X, Y)  # use func.grad not Tensor.grad, since may also use jvp etc.
assert isinstance(Z, torch.Tensor)
assert Z.device.type == "cuda"
</code></pre>

AlphaBetaGamma96: <p>Perhaps <code>torch.func</code> can only handle pytorch primitives, whereas the <code>.grad</code> approach works fine? Perhaps it might be best to get a dev’s opinion: <a class="mention" href="/u/vfdev-5">@vfdev-5</a> (apologies for the tag)</p>
<p>When running the script the following error message I get is,</p>
<pre><code class="lang-auto">torch._dynamo.exc.Unsupported: torch.func.grad(fn) requires the function to be inlined by dynamo

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information
</code></pre>

Lezcano: <p>To be able to use <code>torch.func</code> over NumPy, or compiled functions in general, the relevant <code>torch.func</code> calls should also be within <code>torch.compile</code>.</p>
<p>As a side note, <code>fullgraph=True</code> does not really affect this whole process, it has no side-effects. It’s there to make sure we just get one graph. <code>wrap_numpy</code>, is to be used within <code>torch.compile</code>.</p>
<p>Taking all this into account, we get this script, which does compile and does what you’d expect:</p>
<pre data-code-wrap="python"><code class="lang-python">import numpy as np
import torch
from torch.func import grad

@torch.compiler.wrap_numpy
def fn(x):
    return np.power(x, 2).sum()

@torch.compile
def my_grad(x):
    return torch.func.grad(fn)(x)

arr = torch.tensor([1., 2., 3.], requires_grad=True)
print(my_grad(arr))
</code></pre>
<p>The snipped from the docs from <a href="https://discuss.pytorch.org/t/can-i-use-grad-around-a-numpy-code-and-compile-it/203040/9" class="inline-onebox">Can I use grad around a numpy code and compile it? - #9 by Roy-Kid</a> should work. I guess we had a small regression there. I’ll submit a fix today or tomorrow.</p>

Roy-Kid: <p><a class="mention" href="/u/lezcano">@Lezcano</a> <a class="mention" href="/u/alphabetagamma96">@AlphaBetaGamma96</a> Thanks a lot for your patient help! I have always thought this strategy works, but only a slight bug exists. I have a lot of numpy energy calculators, and I want to reuse them in the NNP package, which is why I am asking this ridiculous question. After it is fixed, I will try again!</p>

Lezcano: <p>It turns out that the example was just wrong. I was calling <code>backward()</code> on a tensor, not a scalar. I sent a fix at <a href="https://github.com/pytorch/pytorch/pull/126872" class="inline-onebox" rel="noopener nofollow ugc">[Docs] Fix NumPy + backward example by lezcano · Pull Request #126872 · pytorch/pytorch · GitHub</a>.<br>
torch.compile already works as advertised to compute gradients in NumPy code.</p>
<p>Note that the usual caveats of minimising graph breaks and avoiding explicit loops / data dependent stuff apply.</p>