Gradcheck for multiple samples

AlphaBetaGamma96: <p><strong>TL;DR</strong> Is there anyway to get <code>gradcheck</code> to work for multiple samples?</p>
<p>Hi All,</p>
<p>I’ve been using <code>torch.autograd.gradcheck</code> to validate a custom <code>torch.autograd.Function</code> object I’ve written, however, I’d like to test the accuracy of the backward method for multiple samples via finite-difference.</p>
<p>The <code>torch.autograd.gradcheck</code> method checks the backward formula analytically for a single sample, but doesn’t work for multiple sample (Or it tries to compute the jacobian across samples, which is something I don’t want).</p>
<p>I did try using <code>torch.func.vmap</code>, however, if one sample fails it’ll raise an error for that one sample, which kind of renders checking multiple samples invalid.</p>
<p>Is there any work around? I did check by using a pytorch equivalent with an automatically computed backward via autograd, and did an <code>torch.allclose</code> call to check element-wise, but that doesn’t get an ‘analytical’ results via finite-difference. Is there any solution to this?</p>

ptrblck: <aside class="quote no-group" data-username="AlphaBetaGamma96" data-post="1" data-topic="179988">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/3da27b/48.png" class="avatar"> AlphaBetaGamma96:</div>
<blockquote>
<p>I did try using <code>torch.func.vmap</code>, however, if one sample fails it’ll raise an error for that one sample, which kind of renders checking multiple samples invalid.</p>
</blockquote>
</aside>
<p>Would using <code>raise_exception=False</code> work as I would assume it wouldn’t break your <code>vmap</code> code?</p>

AlphaBetaGamma96: <p>So, I’ve just tried this and I got an error.</p>
<pre><code class="lang-auto">def custom_func(x):
  return x**3 #but with custom backward

torch.func.vmap(torch.autograd.gradcheck(func=custom_func, inputs=(x), raise_exception=True), in_dims=(0))(x) #fails
</code></pre>
<p>Stacktrace:</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "~/gradcheck_vmap.py", line 126, in &lt;module&gt;
    result = torch.func.vmap(torch.autograd.gradcheck(func=custom_func, inputs=(x), raise_exception=False), in_dims=(0))(x)
  File "~/myenv/lib/python3.10/site-packages/torch/_functorch/vmap.py", line 434, in wrapped
    return _flat_vmap(
  File "~/myenv/lib/python3.10/site-packages/torch/_functorch/vmap.py", line 39, in fn
    return f(*args, **kwargs)
  File "~/myenv/lib/python3.10/site-packages/torch/_functorch/vmap.py", line 619, in _flat_vmap
    batched_outputs = func(*batched_inputs, **kwargs)
TypeError: 'bool' object is not callable
</code></pre>
<p>EDIT: The error is because I pass <code>(x)</code> to the gradcheck function.</p>

ptrblck: <p>You are already executing the <code>gradcheck</code> inside <code>vmap</code> which will directly return a <code>bool</code> value while a function is expected.<br>
Besides that also tensors should be returned so more changes are needed.<br>
Also, I’m using an internal attribute to unwrap the tensor inside <code>vmap</code>, which I assume could easily break, but this code seems to work for now:</p>
<pre><code class="lang-python">import torch
from functools import partial

def custom_func(x):
  return x**3 #but with custom backward

def my_gradcheck(x, func):
    vx = torch._C._functorch.get_unwrapped(x)
    ret = torch.autograd.gradcheck(func=func, inputs=(vx), raise_exception=False)
    return torch.tensor(ret)

x = torch.randn(10, 10, dtype=torch.double, requires_grad=True)

partial_gradcheck = partial(my_gradcheck, func=custom_func)
vmap_gradcheck = torch.func.vmap(partial_gradcheck)
vmap_gradcheck(x)
# tensor([True, True, True, True, True, True, True, True, True, True])
</code></pre>

AlphaBetaGamma96: <p>That’s a neat trick! Thanks for the example <a class="mention" href="/u/ptrblck">@ptrblck</a>!</p>