How to compute the gradient of a component of a vector-valued function?

tmaric: <p>Let’s say I have a function Psi with a 4-dimensional vector output, that takes a 3-dimensional vector u as input. I would like to compute the gradient of the first three components of Psi w.r.t. the respective three components of u:</p>
<pre><code class="lang-auto">import torch

u = torch.tensor([1.,2.,3.], requires_grad=True)

psi = torch.zeros(4)
psi[0] = 2*u[0]
psi[1] = 2*u[1]
psi[2] = 2*u[2]
psi[3] = torch.dot(u,u)

grad_Psi_0 = torch.autograd.grad(psi[0], u[0])
grad_Psi_1 = torch.autograd.grad(psi[1], u[1])
grad_Psi_2 = torch.autograd.grad(psi[2], u[2])
</code></pre>
<p>And I get the error that u[0],u[1], and u[2] are not used in the graph:</p>
<pre><code class="lang-auto">---&gt; 19 grad_Psi_0 = torch.autograd.grad(psi[0], u[0])
     20 grad_Psi_1 = torch.autograd.grad(psi[1], u[1])
     21 grad_Psi_2 = torch.autograd.grad(psi[2], u[2])

File ~/.local/lib/python3.10/site-packages/torch/autograd/__init__.py:275, in grad(outputs, inputs, grad_outputs, retain_graph, create_graph, only_inputs, allow_unused, is_grads_batched)
    273     return _vmap_internals._vmap(vjp, 0, 0, allow_none_pass_through=True)(grad_outputs)
    274 else:
--&gt; 275     return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
    276         outputs, grad_outputs_, retain_graph, create_graph, inputs,
    277         allow_unused, accumulate_grad=False)

RuntimeError: One of the differentiated Tensors appears to not have been used in the graph. Set allow_unused=True if this is the desired behavior.
</code></pre>
<p>Why is this the case? I am using u[i] to build psi…</p>

AlphaBetaGamma96: <p>Your leaf Tensor here is <code>u</code> not <code>u[0]</code> or <code>u[1]</code> so that’s why you have the error you do.</p>
<p>The code you want is shown below</p>
<pre><code class="lang-auto">def func(u):
  return torch.cat([2*u, torch.dot(u,u).unsqueeze(0)], dim=0) #vectorize form of your function

out = torch.autograd.functional.jacobian(func, u)
print(out)
#returns 
tensor([[2., 0., 0.],
        [0., 2., 0.],
        [0., 0., 2.],
        [2., 4., 6.]])
</code></pre>

tmaric: <p>I got a tip on <a href="https://stackoverflow.com/q/73133352/704028" rel="noopener nofollow ugc">Stack Overflow</a> regarding this issue also, so this is what I made work using the answer from there:</p>
<pre><code class="lang-auto">import torch

# u = grad(Phi) = [2*u0, 2*u1, 2*u2]
# Phi = u0**2 + u1**2 + u2**2 = dot(u,u)

u = torch.tensor([1.,2.,3.], requires_grad=True)

psi = torch.zeros(4)
psi[0] = 2*u[0]
psi[1] = 2*u[1]
psi[2] = 2*u[2]
psi[3] = torch.dot(u,u)

print("u = ",u)
print("psi = ",psi)

grad_v_x = torch.autograd.grad(psi[0], u, retain_graph=True)[0]
print(grad_v_x)
grad_v_y = torch.autograd.grad(psi[1], u, retain_graph=True)[0]
print(grad_v_y)
grad_v_z = torch.autograd.grad(psi[2], u, retain_graph=True)[0]
print(grad_v_z)
div_v = grad_v_x[0] + grad_v_y[1] + grad_v_z[2]

# Divergence of the vector phi[0:3]=2u0 + 2u1 + 2u2 w.r.t [u0,u1,u2] = 2+2+2=6
print (div_v)

# laplace(psi[3]) = \partial_u0^2 psi[3] + \partial_u1^2 psi[3] + \partial_u2^2 psi[3]
# = \partial_u0 2x + \partial_u1 2u1 + \partial_u2 2u2 = 2 + 2 + 2 = 6
d_phi_du = torch.autograd.grad(psi[3], u, create_graph=True, retain_graph=True)[0]
print(d_phi_du)
dd_phi_d2u0 = torch.autograd.grad(d_phi_du[0], u, retain_graph=True)[0]
dd_phi_d2u1 = torch.autograd.grad(d_phi_du[1], u, retain_graph=True)[0]
dd_phi_d2u2 = torch.autograd.grad(d_phi_du[2], u, retain_graph=True)[0]

laplace_phi = torch.dot(dd_phi_d2u0 + dd_phi_d2u1 + dd_phi_d2u2, torch.ones(3))

print(laplace_phi)
</code></pre>
<p>The answer from <a class="mention" href="/u/alphabetagamma96">@AlphaBetaGamma96</a> uses the Jacobian, and this variant above uses partial derivative w.r.t. the whole u vector - is there not an option to derive w.r.t. u-component? I mean, with a Jacobian, gradients are computed that are not needed, same as with grad(psi[i], u)</p>

AlphaBetaGamma96: <p>When you calculate gradients via <code>torch.autograd.grad</code> what you’re doing is taking an input vector <code>u</code> (<strong>not</strong> u[0] or u[1] as these are just views on <code>u</code>), you take <code>u</code> and pass it through a model to calculate <code>psi</code>. PyTorch will see entirely <code>u</code> and how it maps to <code>psi</code>. When you call <code>torch.autograd.grad(psi, u)</code> you are telling pytorch to calculate the gradient (via reverse-mode AD) what is the gradient of <code>psi</code> with respect to <code>u</code>. And, as <code>u</code> and <code>psi</code> were used in your model it can quite happily calculate that gradient.</p>
<p><em>however</em> when you pass something like <code>torch.autograd.gradd(psi, u[0])</code> you are saying to PyTorch’s autograd find the gradient of <code>psi</code> and <code>u[0]</code>, but the variable <code>u[0]</code> was never used in the computation of <code>psi</code> the whole vector <code>u</code> was used. When you pass <code>u[0]</code> to <code>torch.autograd.grad</code> you are sending a view on <code>u</code> which <strong>wasn’t</strong> used in the gradient computation.</p>
<p>That’s why autograd gives you an error messages saying,</p>
<pre><code class="lang-auto">RuntimeError: One of the differentiated Tensors appears to not have been used in the graph. Set allow_unused=True if this is the desired behavior.
</code></pre>
<p>because your model used <code>u</code> to compute <code>psi</code>, and not the view of <code>u[0]</code> to compute it. That’s why you get the error you do. Does that make sense?</p>

tmaric: <p>If I understand you correctly, this would mean the operations in constructing psi</p>
<pre><code class="lang-auto">psi = torch.zeros(4)
psi[0] = 2*u[0]
psi[1] = 2*u[1]
psi[2] = 2*u[2]
psi[3] = torch.dot(u,u)
</code></pre>
<p>that use components of u, do not propagate the information to the graph, that those same u-components take part in building psi-components? If this is true, I find this kinda surprising. If u is a vector and I build new tensors using it’s components explicitly like this, and taking a component (slicing a tensor) is an operation that can be differentiated (?), I would expect the dependency information of psi[i] → depends → u[i] would be stored in the graph.</p>
<p>Then computing grad(psi, u[i]) would make sense, because mathematically, I am taking a Jacobian of a vector-valued function w.r.t it’s one leaf u[i] input variable… on paper it makes sense, at least to me.</p>

AlphaBetaGamma96: <aside class="quote no-group" data-username="tmaric" data-post="5" data-topic="157581">
<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/t/f19dbf/48.png" class="avatar"> tmaric:</div>
<blockquote>
<p>I find this kinda surprising. If u is a vector and I build new tensors using it’s components explicitly like this</p>
</blockquote>
</aside>
<p>But the thing is, you’re not building a new tensor via <code>u[0]</code> you’re taking a view on <code>u</code> for its 0-th element. That’s why autograd can’t see it how individual components of <code>u</code> act within the computation graph it only sees <code>u</code> not <code>u[0]</code> nor <code>u[1]</code>. That’s why you get that error about unused Tensors.</p>
<pre><code class="lang-auto">u = torch.randn(3, requires_grad=True)
print(u[0]) #returns tensor(-0.0611, grad_fn=&lt;SelectBackward0&gt;)
</code></pre>

AlphaBetaGamma96: <p>Think of this computation via this (pretty poor ASCII diagram)</p>
<pre><code class="lang-auto">u -&gt; model -&gt; psi 
</code></pre>
<p>When you do <code>torch.autograd.grad(psi, u[0])</code></p>
<pre><code class="lang-auto">u -&gt; model -&gt; psi

|
v

u[0] 
</code></pre>
<p>As you can see, <code>psi</code> and <code>u[0]</code> have no connection and <em>hence</em> no gradient</p>

tmaric: <p><img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji only-emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
<blockquote>
<p>But the thing is, you’re not building a new tensor via <code>u[0]</code> you’re taking a view on <code>u</code> for its 0-th element. That’s why autograd can’t see it how individual components of <code>u</code> act within the computation graph it only sees <code>u</code> not <code>u[0]</code> nor <code>u[1]</code> . That’s why you get that error about unused Tensors.</p>
</blockquote>
<p>Yeah, thanks a lot! I think I got that, I’m just saying, as a person just starting to use torch.autograd, with another background (where there are no views), just after reading about the DAG/AD, etc, I wasn’t expecting this to be the case. Maybe it would be good to add this small example to the official documentation and warn newcomers like myself about a[0] not reaching into u, but generating a new view, that’s then not accounted for in the DAG.</p>
<p>If I look at how my psi is defined:</p>
<pre><code class="lang-auto">psi = torch.zeros(4)
psi[0] = 2*u[0]
psi[1] = 2*u[1]
psi[2] = 2*u[2]
psi[3] = torch.dot(u,u)
</code></pre>
<p>mathematically, its components are constructed as functions of the components of the vector u, so to my beginner’s mind,</p>
<p><code>u[0], u[1], u[2] -&gt; psi</code></p>
<p>or</p>
<p><code>psi = psi(u[0],u[1],u[2])</code></p>
<p>then I should be able to do d psi / du[2] in this graph. And what you’re saying (if I understood you), when I call u[0], I create a <em>separate</em> view of the component of u, disconnected from the way psi is computed? Maybe I am not understanding it after all…</p>

AlphaBetaGamma96: <aside class="quote no-group" data-username="tmaric" data-post="8" data-topic="157581">
<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/t/f19dbf/48.png" class="avatar"> tmaric:</div>
<blockquote>
<p>I call u[0], I create a <em>separate</em> view of the component of u, disconnected from the way psi is computed?</p>
</blockquote>
</aside>
<p>Exactly, when you call <code>torch.autograd.grad(psi, u[0])</code> you are telling autograd to find the gradient of <code>psi</code> with respect to <code>u[0]</code>. However, <code>psi</code> never used <code>u[0]</code> to be calculated only <code>u</code> was used, and hence there’s no gradient.</p>