Get all zero answer while calculating jacobian in PyTorch using build-in function-jacobian

miraboreasu: <p>I am trying to compute Jacobian matrix, it is computed between two vectors, and the result should be a matrix. Ref:</p><aside class="onebox allowlistedgeneric" data-onebox-src="https://www.geeksforgeeks.org/jacobian-matrix-in-pytorch/">
  <header class="source">
      <img src="https://media.geeksforgeeks.org/wp-content/cdn-uploads/gfg_favicon.png" class="site-icon" width="16" height="16">

      <a href="https://www.geeksforgeeks.org/jacobian-matrix-in-pytorch/" target="_blank" rel="noopener nofollow ugc" title="03:28AM - 27 May 2021">GeeksforGeeks – 27 May 21</a>
  </header>

  <article class="onebox-body">
    <img src="https://media.geeksforgeeks.org/wp-content/cdn-uploads/gfg_200x200-min.png" class="thumbnail onebox-avatar" width="200" height="200">

<h3><a href="https://www.geeksforgeeks.org/jacobian-matrix-in-pytorch/" target="_blank" rel="noopener nofollow ugc">Jacobian matrix in PyTorch - GeeksforGeeks</a></h3>

  <p>A Computer Science portal for geeks. It contains well written, well thought and well explained computer science and programming articles, quizzes and practice/competitive programming/company interview Questions.</p>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<pre><code class="lang-auto">import torch
from torch.autograd.functional import jacobian
from torch import tensor

def get_f (x):      
    f=torch.arange(0,3, requires_grad=True, dtype=torch.float64)
    for i in looparray:
        with torch.no_grad():   
            f[i] = x[i]**2    
    return f
    
looparray=torch.arange(0,3)
x=torch.arange(0,3, requires_grad=True, dtype=torch.float64)
J = jacobian(get_f, x).detach().numpy()
</code></pre>
<p>Mathematically, Jacobian is computed by</p>
<pre><code class="lang-auto">df[0]/dx[0]   df[0]/dx[1]  df[0]/dx[2]
df[1]/dx[0]   df[1]/dx[1]  df[1]/dx[2]
df[2]/dx[0]   df[2]/dx[1]  df[2]/dx[2]
</code></pre>
<p>It should be</p>
<pre><code class="lang-auto">0 0 0
0 2 0
0 0 4
</code></pre>
<p>But the result from the code is</p>
<pre><code class="lang-auto">0 0 0
0 0 0
0 0 0
</code></pre>
<p>And I must update in-place in the <code>get_f</code>, like <code>f</code> maybe becomes</p>
<pre><code class="lang-auto"> f[0] = x[0]**2
 f[1] = x[0]**2+x[1]**2
 f[2] = x[1]**2+x[2]**2
</code></pre>
<p>So a for loop (here 3x3 is just an example, it could be way larger) or in-place updates is very necessary in my case</p>

AlphaBetaGamma96: <aside class="quote no-group" data-username="miraboreasu" data-post="1" data-topic="158340">
<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/m/6bbea6/48.png" class="avatar"> miraboreasu:</div>
<blockquote>
<pre><code class="lang-auto">        with torch.no_grad():   
</code></pre>
</blockquote>
</aside>
<p>Remove this line. It’s resulting in no gradients flowing through your function <code>get_f</code></p>

miraboreasu: <p>Hey, no it won’t, it just track this indexing operation. And if I remove it, error like this pop yp<br>
RuntimeError: a view of a leaf Variable that requires grad is being used in an in-place operation.</p>

AlphaBetaGamma96: <p>It’s best not to define arrays and update in-place within the function, the following code works,</p>
<pre><code class="lang-auto">import torch
from torch.autograd.functional import jacobian
from torch import tensor

def get_f(x):
  return x**2

x=torch.arange(0,3, requires_grad=True, dtype=torch.float64)
J = jacobian(get_f, x).detach().numpy()
print(J)
</code></pre>

miraboreasu: <p>Thank you, but I must update in-place, any idea?</p>

AlphaBetaGamma96: <p>Autograd doesn’t work with in-place updates. How would you define the gradient of placing a scalar within an array?</p>

miraboreasu: <p>Thank you, but are you sure that no way to work with in-place updates? I remember I did one before, but if I want to have some complicate f, different updates for different indexes, how can I do it</p>

AlphaBetaGamma96: <p>Why do you need to do it via in-place updates? Why can’t you just return a Tensor that the network creates itself? Do you have an example to share?</p>

miraboreasu: <p>Like if <code>f</code> becomes</p>
<pre><code class="lang-auto"> f[0] = x[0]**2
 f[1] = x[0]**2+x[1]**2
 f[2] = x[1]**2+x[2]**2
</code></pre>

AlphaBetaGamma96: <p>why don’t you just define a function like,</p>
<pre><code class="lang-auto">def f(x,y,z):
  return x**2, x**2+y**2, y**2+z**2
</code></pre>
<p>(or the equivalent index version)?</p>

miraboreasu: <p>Thank you, but here is just an example 3 by 3. Actually, in my case, I may have hundreds of variables, and that is why I use for loop.</p>

AlphaBetaGamma96: <p>If you’re doing it for multiple variables you could you could forgo pre-defining <code>f</code> like above</p>
<aside class="quote no-group" data-username="miraboreasu" data-post="1" data-topic="158340">
<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/m/6bbea6/48.png" class="avatar"> miraboreasu:</div>
<blockquote>
<pre><code class="lang-auto">    f=torch.arange(0,3, requires_grad=True, dtype=torch.float64)
</code></pre>
</blockquote>
</aside>
<p>and use <code>torch.stack</code> and list comprehension to get AD to work. So, for example,</p>
<pre><code class="lang-auto">x = torch.arange(1,5,dtype=torch.float32, requires_grad=True)
out = torch.stack([ x[i]**2 for i in range(x.shape[0])], 0)
print(out.grad_fn) #returns &lt;StackBackward0 object at 0x7f3415b34610&gt;
</code></pre>

miraboreasu: <p>I am sorry but I didn’t get it. It is multi variables, and <code>f</code> is also multiple</p>

AlphaBetaGamma96: <p>One thing I don’t understand is why the updates need to be in-place? With your example above,</p>
<aside class="quote no-group" data-username="miraboreasu" data-post="1" data-topic="158340">
<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/m/6bbea6/48.png" class="avatar"> miraboreasu:</div>
<blockquote>
<pre><code class="lang-auto"> f[0] = x[0]**2
 f[1] = x[0]**2+x[1]**2
 f[2] = x[1]**2+x[2]**2
</code></pre>
</blockquote>
</aside>
<p>you can try something like this,</p>
<pre><code class="lang-auto">import torch
from torch.autograd.functional import jacobian
from torch import tensor
"""
#old function with f pre-allocated
def get_f (x):      
  f=torch.arange(0,3, requires_grad=True, dtype=torch.float64)
  for i in looparray:
    f[i] = x[i]**2    
  return f
"""
#new func
def get_f(x):
  f0 = x[0]**2
  f1 = x[0]**2+x[1]**2
  f2 = x[1]**2+x[2]**2
  return torch.stack([f0, f1, f2], dim=0)

x=torch.arange(0,3, requires_grad=True, dtype=torch.float64)
J = jacobian(get_f, x).detach().numpy()

print(J) #returns 
#[[0. 0. 0.]
# [0. 2. 0.]
# [0. 2. 4.]]
</code></pre>
<p>When you calculate gradients via Automatic Differentiation (AD) you need to construct the output because AD <em>uses</em> the output Tensor to calculate the gradients which is why I don’t think you can use in-place updates. When you have a constant Tensor I don’t think AD can track what the derivative would be. But it’d be best to get a dev to confirm the exact details.</p>

miraboreasu: <p>Now I get your idea, but I still have problem. the reason for the in-place update is that I may have many <code>f</code>, I cannot do <code>f1=</code>, <code>f2=</code>, <code>f3=</code> to <code>f100=</code>, also, I need <code>+=</code> I send you the exact example via message, since it is too long for here.</p>
<p>That is why I use <code>for</code> (in-place or index) to construct <code>f</code></p>

AlphaBetaGamma96: <p>If it’s too long to implement directly you could calculate each value and append each value into a list then stack into a single Tensor (like the example above),. Also, make sure to avoid use of <code>+=</code> or <code>*=</code> as these are in-place operations (and are not support by AD, you need to use the out-of-place equivalent) and don’t use <code>torch.no_grad()</code> as that stops AD tracking gradients.</p>
<p>For reference, the out-of-place is,</p>
<pre><code class="lang-auto">x *= 5 #in-place
x  = x * 5 #out-of-place
</code></pre>
<pre><code class="lang-auto">values = []
for i in range(3):
  val = calc_value(i) #your function for a given index
  values.append(val) #append to list
values = torch.stack(values) #converts to Tensor, do NOT use torch.tensor()
</code></pre>
<p>Make sure to use <code>torch.stack</code> instead of <code>torch.tensor</code> as <code>torch.tensor</code> will create a new Tensor and free the computational graph for your values which will give no gradient.</p>

miraboreasu: <p>I see, instead of given a whole <code>f</code>, “push back” each element to <code>f</code></p>

AlphaBetaGamma96: <p>But don’t pre-allocate the memory for <code>f</code>, you’ll need to create an empty list each time in order to get derivatives.</p>

miraboreasu: <p>I just found that in my case, the <code>f</code> sometimes is not computed in order, some data could be added and revisit to the previous locations. What I can do with it</p>

AlphaBetaGamma96: <p>I don’t know why it’s giving different permutations of the output but that’s something you can try and solve with <code>permutations</code> or make your function deterministic instead. Good luck!</p>