In what situation does the decorator once_differentiable help?

Fei_Liu: <p>I’m reading the <a href="https://pytorch.org/docs/stable/notes/extending.html#extending-torch-autograd" rel="noopener nofollow ugc">Extending <code>torch.autograd</code></a> doc, and the Double Backward with Custom Functions <a href="https://pytorch.org/tutorials/intermediate/custom_function_double_backward_tutorial.html" rel="noopener nofollow ugc">tutorial</a>. I think I have a grasp of the idea and why it matters for higher order gradient computations. That said, I have a feeling that <code>once_differentiable</code> provides very limited helps in terms of debugging and error prompting. Consider this toy <code>x ** 2</code> example (a simplified version of the <a href="https://pytorch.org/tutorials/intermediate/custom_function_double_backward_tutorial.html#saving-intermediate-results-what-not-to-do" rel="noopener nofollow ugc">not-to-do example</a> from the tutorial):</p>
<pre><code class="lang-auto">import torch

class BadSquare(torch.autograd.Function):
    # This is an example of what NOT to do!
    @staticmethod
    def forward(ctx, x):
        x2 = x * 2
        ctx.save_for_backward(x2)
        return x ** 2

    @staticmethod
    def backward(ctx, grad_out):
        x2, = ctx.saved_tensors
        return grad_out * x2

x = torch.rand(2, requires_grad=True, dtype=torch.double)
torch.autograd.gradcheck(BadSquare.apply, x)  # first order check passes
torch.autograd.gradgradcheck(BadSquare.apply, x)  # fails, grad(grad(x)) Jacobian mismatch
</code></pre>
<p>Instead of ctx saving <code>x</code> and using it in the backward call, it creates an intermediate tensor <code>x2</code> which is detached from the computational graph (in the no-grad forward context) and use that in backward (bad!), making <code>x.grad</code> independent from <code>x</code>. Hence the grad(grad(x)) check failed.</p>
<p>Now, since this <code>BadSquare</code> function doesn’t support double backward, I mark the backward as once_differentiable:</p>
<pre><code class="lang-auto">    @staticmethod
    @torch.autograd.function.once_differentiable  # add this
    def backward(ctx, grad_out):
        x2, = ctx.saved_tensors
        return grad_out * x2
</code></pre>
<p>However the only thing this decorator does is to check whether <code>grad_out</code> requires_grad in this call (it’ll error if yes). The thing is, for almost all of such “bad examples”, the key limitation is that we have some intermediate results that are hard (or impossible) to be associated to the input x via the standard computational graph (e.g. we may need some third-party approach to compute the grad). However, this has nothing to do with <code>grad_out</code> and it barely helps to know whether grad_out needs grad or not (<strong>Side question: is there even a case where grad_out may require grad?</strong> ).</p>
<p>If I compute the Hessian:</p>
<pre><code class="lang-auto">def bad_sum_squared(x: torch.Tensor):
    return BadSquare.apply(x).sum()

x = torch.rand(2, requires_grad=True, dtype=torch.double)
h = torch.autograd.functional.hessian(bad_sum_squared, x)
print(h)  # all zeros
</code></pre>
<p>It gives all zeros <strong>both with and without</strong> the <code>once_differentiable</code> decorator. There’s no error prompted with the decorator. It just gives a wrong result. So to me this behaves like we already have a backward that doesn’t support higher order derivatives – the computed results are considered as “const” w.r.t. the inputs (like linear/affine functions). Adding the decorator does nothing more than an emphasis that “<strong>we gave up</strong>”. It doesn’t produce errors as intended.</p>
<p>What am I missing here?</p>

Fei_Liu: <p>Another post mentioning <code>once_differentiable</code>:</p><aside class="quote" data-post="1" data-topic="31513">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/jxgu1016/48/1153_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/what-does-the-function-wrapper-once-differentiable-do/31513">What does the function wrapper @once_differentiable do?</a> <a class="badge-category__wrapper " href="/c/autograd/7"><span data-category-id="7" style="--category-badge-color: #AB9364; --category-badge-text-color: #FFFFFF;" data-drop-close="true" class="badge-category " title="A category of posts relating to the autograd engine itself."><span class="badge-category__name">autograd</span></span></a>
  </div>
  <blockquote>
    I find some codes in <a href="https://github.com/facebookresearch/maskrcnn-benchmark/blob/55d3ab44aa9ff5261afb8244ab5e265266d8436f/maskrcnn_benchmark/layers/roi_align.py#L25" rel="noopener nofollow ugc">this</a> and become curious about the function of @once_differentiable.
  </blockquote>
</aside>

<p>We say that it prevents higher order derivatives computation, but like I mentioned above, I doubt if it makes any difference if we just remove the decorator from the example code in the post – in either case, the higher order derivatives are still “computable”, it’s just all zeros (and again, with or without the decorator).</p>

soulitzer: <p>If you had a function for which higher derivatives exist, but you silently return incorrect results (e.g. zero), then once_differentiable could be a good way to explicitly raise an error when someone tries to access higher-order derivatives.</p>

Fei_Liu: <p>The thing is, if someone tries to access higher-order derivatives via double backward, then it must be that there’s a grad_output in the backward param list that requires grad. However, that’s often not the case when someone tries to encapsulate fancy/complicated logics in the forward/backward implementation, where the only params in the backward call are anyways irrelevant to higher order derivative computations. This makes the <code>once_differentiable</code> decorator pretty much useless. E.g. I have given an example in the main post that it doesn’t stop wrong Hessian computation. Another example is the <a href="https://discuss.pytorch.org/t/what-does-the-function-wrapper-once-differentiable-do/31513">post</a> above, which I also believe it does nothing, because the only access exposed to the outside is grad_output, which are always empty in higher order derivative computations.</p>