How does one reuse the autograd computational graph

bradbell: <p>I tried posting this as a continuation of the following post and did not get a response:</p><aside class="quote" data-post="4" data-topic="87293">
  <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/alband/48/215_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/reusing-jacobian-and-hessian-computational-graph/87293/4">Reusing Jacobian and Hessian computational graph</a> <a class="badge-category__wrapper " href="/c/autograd/7"><span data-category-id="7" 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>
    It seems that PyTorch is more or less at the level as JAX for the gradient computation 

That’s good news, that needs that your ops are big enough that the creation of the graph is negligible. 

Hessian is about 20x slower in PyTorch than in JAX. 

That might be due to other reasons, mainly the way the Hessian is computed: jax can use forward mode AD to speed this up while pytorch does not have forward mode AD (yet <img width="20" height="20" src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=12" title="wink" alt="wink" class="emoji"> ).
  </blockquote>
</aside>

<p>Consider the example case of computing determinants using expansion by minors.<br>
I know this is not an efficient way to compute determinants, but it demonstrates a case that has lots of operations on scalar values.</p>
<p>I am trying to re-use the graph for this case and my method of doing so works for 2 by 2 matrices and fails for 3 by 3 matrices;  see the pythong programs below, What am I doing wrong ?</p>
<p>det_22.py</p>
<pre><code class="lang-auto"># Test reusing graph for derivatives of determinant of 2 by 2 matrix.
# The output generated by this program is below:
#
# First gradient passed check.
# Second gradient passed check.
#
# imports
import sys
import torch
import numpy
#
# check_grad
def check_grad(ax) :
   # ok, eps99
   ok    = True
   eps99 = 99.0 * numpy.finfo(float).eps
   #
   # ok
   check = ax.data[1,1]
   if abs( ax.grad[0,0] - check ) &gt; eps99 :
      ok = False
      print( f'ax.grad[0,0] = {ax.grad[0,0]}, check = {check}' )
   #
   # ok
   check = -ax.data[1,0]
   if abs( ax.grad[0,1] - check ) &gt; eps99 :
      ok = False
      print( f'ax.grad[0,1] = {ax.grad[0,1]}, check = {check}' )
   #
   # ok
   check = -ax.data[0,1]
   if abs( ax.grad[1,0] - check ) &gt; eps99 :
      ok = False
      print( f'ax.grad[1,0] = {ax.grad[1,0]}, check = {check}' )
   #
   # ok
   check = ax.data[0,0]
   if abs( ax.grad[1,1] - check ) &gt; eps99 :
      ok = False
      print( f'ax.grad[1,1] = {ax.grad[1,1]}, check = {check}' )
   #
   return ok
#
# main
def main() :
   #
   # ok
   ok = True
   #
   # n
   n = 2
   #
   # ax
   x  = numpy.random.uniform(0.0, 1.0, (n , n) )
   ax = torch.tensor(x, requires_grad = True)
   #
   # az
   az  = ax[0,0] * ax[1,1] - ax[0,1] * ax[1,0]
   #
   # ax.grad
   az.backward(retain_graph = True)
   #
   # check_grad
   if check_grad(ax) :
      print( 'First gradient passed check.' )
   else :
      print( 'First gradient failed check.' )
   #
   #
   # ax.data
   x  = numpy.random.uniform(0.0, 1.0, (n, n) )
   for i in range(n) :
      for j in range(n) :
         ax.data[i,j] = x[i,j]
   #
   # ax.grad
   ax.grad.zero_()
   az.backward(retain_graph = True)
   #
   # check_grad
   if check_grad(ax) :
      print( 'Second gradient passed check.' )
   else :
      print( 'Second gradient failed check.' )
#
main()
</code></pre>
<p>det_33.py</p>
<pre><code class="lang-auto"># Test reusing graph for derivatives of determinant of 3 by 3 matrix.
# The output generated by this program is below. The actual numbers
# in the output will vary because a different random matrix is chosen
# for each evaluation.
#
# First gradient passed check.
# ax.grad[0,0] = 0.07585514040844837, check = 0.4295608074373773
# ax.grad[0,1] = -0.6133183512861293, check = -0.11782369019260797
# ax.grad[0,2] = 0.5337097801031835, check = 0.040633019648616306
# Second gradient failed check.
#
#
# imports
import torch
import numpy
#
# check_grad
def check_grad(ax) :
   # ok, eps99
   ok    = True
   eps99 = 99.0 * numpy.finfo(float).eps
   #
   # ok
   check = ( ax[1,1] * ax[2,2] - ax[1,2] * ax[2,1] )
   if abs( ax.grad[0,0] - check ) &gt; eps99 :
      ok = False
      print( f'ax.grad[0,0] = {ax.grad[0,0]}, check = {check}' )
   #
   # ok
   check = - ( ax[1,0] * ax[2,2] - ax[1,2] * ax[2,0] )
   if abs( ax.grad[0,1] - check ) &gt; eps99 :
      ok = False
      print( f'ax.grad[0,1] = {ax.grad[0,1]}, check = {check}' )
   #
   # ok
   check = ( ax[1,0] * ax[2,1] - ax[1,1] * ax[2,0] )
   if abs( ax.grad[0,2] - check ) &gt; eps99 :
      ok = False
      print( f'ax.grad[0,2] = {ax.grad[0,2]}, check = {check}' )
   #
   return ok
#
# main
def main() :
   #
   # ok
   ok = True
   #
   # n
   n = 3
   #
   # ax
   x  = numpy.random.uniform(0.0, 1.0, (n, n))
   ax = torch.tensor(x, requires_grad = True)
   #
   #  ax[0,0]  ax[0,1]  ax[0,2]
   #  ax[1,0]  ax[1,1]  ax[1,2]
   #  ax[2,0]  ax[2,1]  ax[2,2]
   #
   # az
   az  = ax[0,0] * ( ax[1,1] * ax[2,2] - ax[1,2] * ax[2,1] )
   az -= ax[0,1] * ( ax[1,0] * ax[2,2] - ax[1,2] * ax[2,0] )
   az += ax[0,2] * ( ax[1,0] * ax[2,1] - ax[1,1] * ax[2,0] )
   #
   #
   # ax.grad
   az.backward(retain_graph = True)
   #
   # check_grad
   if check_grad(ax) :
      print( 'First gradient passed check.' )
   else :
      print( 'First gradient failed check.' )
   #
   #
   # ax.data
   x  = numpy.random.uniform(0.0, 1.0, (n, n) )
   for i in range(n) :
      for j in range(n) :
         ax.data[i,j] = x[i,j]
   #
   # ax.grad
   ax.grad.zero_()
   az.backward(retain_graph = True)
   #
   # check_grad
   if check_grad(ax) :
      print( 'Second gradient passed check.' )
   else :
      print( 'Second gradient failed check.' )
#
main()
</code></pre>

KFrank: <p>Hi Brad!</p>
<p>The short story is that you are using <code>.data</code> which is deprecated (in the<br>
public-facing api) and can lead to errors (such as yours).</p>
<aside class="quote no-group" data-username="bradbell" data-post="1" data-topic="190447" data-full="true">
<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/b/e95f7d/48.png" class="avatar"> bradbell:</div>
<blockquote>
<p>I am trying to re-use the graph for this case and my method of doing so works for 2 by 2 matrices</p>
</blockquote>
</aside>
<p>This works for your 2-by-2 case, in essence, by accident.</p>
<aside class="quote no-group">
<blockquote>
<p>and fails for 3 by 3 matrices;<br>
…<br>
What am I doing wrong ?<br>
…<br>
det_33.py</p>
<pre><code class="lang-auto">def main() :
   ...
   az  = ax[0,0] * ( ax[1,1] * ax[2,2] - ax[1,2] * ax[2,1] )
   az -= ax[0,1] * ( ax[1,0] * ax[2,2] - ax[1,2] * ax[2,0] )
   az += ax[0,2] * ( ax[1,0] * ax[2,1] - ax[1,1] * ax[2,0] )
</code></pre>
</blockquote>
</aside>
<p>When you compute <code>az</code>, autograd stores some of the product tensors,<br>
e.g., <code>ax[1,1] * ax[2,2]</code>, as temporary values <em>in the computation graph<br>
itself,</em> for use during the backward pass.</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">   for i in range(n) :
      for j in range(n) :
         ax.data[i,j] = x[i,j]
</code></pre>
</blockquote>
</aside>
<p>By using <code>.data</code> (whose public use is deprecated), you hide the modification<br>
of <code>ax</code> from autograd.  The temporary tensor stored in the computation graph<br>
is now out of date with respect to the modified value of <code>ax</code> (but autograd<br>
doesn’t know this).  So when you call <code>az.backward()</code> the second time,<br>
autograd uses the stale temporary tensor and computes an incorrect gradient.</p>
<p>Consider:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
2.1.0
&gt;&gt;&gt;
&gt;&gt;&gt; a = torch.tensor ([1.0, 2.0, 3.0], requires_grad = True)
&gt;&gt;&gt;
&gt;&gt;&gt; # autograd will save  a[0] * a[1]  as an intermediate result in the computation graph
&gt;&gt;&gt; # let's call it  tmp = tensor(2.0, grad_fn=&lt;MulBackward0&gt;)
&gt;&gt;&gt; # it will be used during p.backward() to compute  a.grad[2] = 2.0
&gt;&gt;&gt; p = a[0] * a[1] * a[2]
&gt;&gt;&gt;
&gt;&gt;&gt; p.backward (retain_graph = True)
&gt;&gt;&gt; a.grad
tensor([6., 3., 2.])
&gt;&gt;&gt;
&gt;&gt;&gt; x = torch.tensor ([10.0, 20.0, 30.0])
&gt;&gt;&gt; for  i in range (3):
...     a.data[i] = x[i]   # using .data tricks autograd
...
&gt;&gt;&gt; a.grad.zero_()
tensor([0., 0., 0.])
&gt;&gt;&gt; p.backward()           # uses stale tmp to compute incorrect value for a.grad[2]
&gt;&gt;&gt; a.grad
tensor([600., 300.,   2.])
&gt;&gt;&gt;
&gt;&gt;&gt; # let's not trick autograd
&gt;&gt;&gt; y = torch.tensor ([2.0, 4.0, 6.0])
&gt;&gt;&gt; for  i in range (3):
...     a[i] = x[i]        # no .data -- autograd works and catches your error
...
Traceback (most recent call last):
  File "&lt;stdin&gt;", line 2, in &lt;module&gt;
RuntimeError: a view of a leaf Variable that requires grad is being used in an in-place operation.
</code></pre>
<p>To be clear, you <em>can’t</em> reuse the computation graph in the way you propose<br>
because, in addition to storing a record of the forward-pass operations, it<br>
also stores various data-dependent intermediate tensors for use in the<br>
backward pass.  When the input data changes, you have to rebuild a new<br>
computation graph (that is, perform the forward pass again), because you<br>
have to compute new values for those data-dependent intermediate tensors.</p>
<p>Best.</p>
<p>K. Frank</p>