How to get separated gradients from torch.autograd

zhangqian-sh: <p>Hi,</p>
<p>I am trying to get the gradients of a function with multiple vector outputs. Let <code>x, y, z</code> as input and <code>u,v,w</code> as output, and they are all vectors with length <code>N</code>. if I use <code>torch.autograd.grad</code>, my code is like</p>
<pre><code class="lang-python">def func(x):
    do something fancy to x

def gradients(outpt, inpt):
    return torch.autograd.grad(
        outputs, inputs, grad_outputs=torch.ones_like(outpt), create_graph=True
    )[0]

inpt = torch.stack((x, y, z), 1)
outpt = func(inpt)
grad = gradients(outpt, inpt)
</code></pre>
<p>Then the <code>grad</code> is a  <code>N*3</code> shape matrix, and <code>grad[:, 0]</code> is actually $u_x+u_y+u_z$ (the summation of partial derivatives of u to x, y and z). But I want to get $u_x, u_y, u_z$ separately. I have tried the <code>torch.autograd.functional.jacobian</code> as following,</p>
<pre><code class="lang-python">def func(x):
    do something fancy to x

def jacobian(f, x):
    return torch.sum(
        torch.autograd.functional.jacobian(f, x, create_graph=True), axis=0
    )

inpt = torch.stack((x, y, z), 1)
grad = jacobian(func, inpt)
</code></pre>
<p>This can give me the <code>grad</code> contains separate gradients of <code>u, v, w</code> to <code>x, y, z</code>, the <code>grad</code> size is <code>3*N*3</code>, where $u_y$ (the derivative of u to y) is <code>grad[0, :, 1]</code>, etc. But it will compute the derivatives for each single output element to each single input element, which takes too much GPU memory unnecessarily.</p>
<p>So, is there an elegant way to get the gradients I want with <code>torch.autograd</code>? Thank you so much!</p>

JosueCom: <p>I have a similar question. Is there an effective way to calculate the derivative of a vector output with respect to a single scalar input?</p>

KFrank: <p>Hi Josue!</p>
<aside class="quote no-group" data-username="JosueCom" data-post="2" data-topic="132754" data-full="true">
<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/josuecom/48/27043_2.png" class="avatar"> JosueCom:</div>
<blockquote>
<p>Is there an effective way to calculate the derivative of a vector output with respect to a single scalar input?</p>
</blockquote>
</aside>
<p>Yes, you may use pytorch’s <a href="https://pytorch.org/docs/1.13/autograd.html#forward-mode-automatic-differentiation" rel="noopener nofollow ugc">forward-mode automatic differentiation.</a></p>
<p>See this post for an example:</p>
<aside class="quote quote-modified" data-post="2" data-topic="164854">
  <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/k/ecb155/48.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/batch-differentiating-with-respect-to-single-parameter/164854/2">Batch-differentiating with respect to single parameter</a> 
  </div>
  <blockquote>
    Hi Evan! 

You can’t eliminate the loop using backward-mode autograd.  Autograd 
will let you compute the derivatives of a single scalar (e.g., a loss) with 
respect to a batch of variables (i.e., compute the gradient) in a single 
pass, but it won’t compute the derivatives of a batch of results with 
respect to a single variable in a single pass. 

Yes (but with the proviso that forward-mode autograd is still in beta / 
experimental). 
I have not used forward-mode autograd for anything real so …
  </blockquote>
</aside>

<p>Best.</p>
<p>K. Frank</p>