Computing derivative of NN output wrt to input

VictorSu33: <p>I’m currently trying to implement Physics Informed Neural Networks to solve ODEs. As part of it, I need to compute the derivative of the NN wrt to the inputs. So far I have used the code</p>
<pre><code class="lang-auto">dy_dt = torch.autograd.grad(y,t, torch.ones_like(y), create_graph=True)[0]
</code></pre>
<p>Where t is some tensor with shape [N,1].<br>
This worked for learning scalar valued functions where the NN had 1 input node and one output node returning a tensor of shape [N,1]. However now I want to try and solve planar systems so my NN should have 1 input node and 2 output nodes. When I create a training tensor which has shape [N,1], my NN output has shape [N,2] which seems fine, but the resulting derivative has shape [N,1]. I also tried to use</p>
<pre><code class="lang-auto">torch.autograd.functional.jacobian(model, t)
</code></pre>
<p>but this returned a tensor of shape [N,2,N,1]. I think it should be possible to access the right derivatives from this but I was wondering if there was a to return a tensor of shape [N,2], so each row is the derivative of the output of a specific input wrt to that input.</p>

KFrank: <p>Hi Victor!</p>
<aside class="quote no-group" data-username="VictorSu33" data-post="1" data-topic="206175" 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/victorsu33/48/71373_2.png" class="avatar"> VictorSu33:</div>
<blockquote>
<p>Where t is some tensor with shape [N,1].</p>
</blockquote>
</aside>
<p>I assume that <code>N</code> is the batch dimension of your input tensor.</p>
<aside class="quote no-group">
<blockquote>
<p>When I create a training tensor which has shape [N,1], my NN output has shape [N,2]</p>
</blockquote>
</aside>
<p>I also assume that <code>N</code> is the batch dimension of your output tensor and<br>
also that your <code>model</code> doesn’t mix batch elements together.  That is, that<br>
<code>model (t)[i]</code> only depends on <code>t[i]</code>.</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">torch.autograd.functional.jacobian(model, t)
</code></pre>
<p>but this returned a tensor of shape [N,2,N,1]. I think it should be possible to access the right derivatives from this</p>
</blockquote>
</aside>
<p>Based on my assumption that your <code>model</code> doesn’t mix together batch<br>
elements, the “batch-off-diagonal” elements of the jacobian will be zero.</p>
<p>In your case, it might be simplest simply to extract the non-zero elements<br>
of the jacobian (which will be the derivatives of your two scalar outputs<br>
with respect to your single scalar input on a per-batch-element basis).</p>
<p>If you do this, however, autograd will be unnecessarily computing the<br>
zero, batch-off-diagonal elements of the full jacobian, which could matter<br>
for your performance if <code>N</code> is large.</p>
<aside class="quote no-group">
<blockquote>
<p>but I was wondering if there was a to return a tensor of shape [N,2], so each row is the derivative of the output of a specific input wrt to that input.</p>
</blockquote>
</aside>
<p>Yes.  Probably the best way would be to apply <code>vmap</code> to <code>jacfwd</code>.*  This will<br>
apply <code>jacfwd</code> to <code>model</code> on a per-batch-element basis, giving your desired<br>
derivatives, without running a python loop (which would dramatically slow<br>
things down).</p>
<p>Here is a script that illustrates these two approaches:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

_ = torch.manual_seed (2024)

t = torch.randn (3, 1)
model = torch.nn.Sequential (torch.nn.Linear (1, 2), torch.nn.Tanh(), torch.nn.Linear (2, 2))

jac = torch.autograd.functional.jacobian (model, t)    # "batch-off-diagonal" elements are zero
print ('jac = ...')
print (jac)

# equivalently, could have used
#   jac = torch.func.jacfwd (model) (t)
# or
#   jac = torch.func.jacrev (model) (t)

grad = jac.squeeze().diagonal (dim1 = 0, dim2 = 2).T   # extract non-zero, "diagonal" elements
print ('grad = ...')
print (grad)

gradB = torch.func.vmap (torch.func.jacfwd (model), in_dims = (0,)) (t).squeeze()
print ('gradB = ...')
print (gradB)

print ('torch.equal (grad, gradB):', torch.equal (grad, gradB))
</code></pre>
<p>And here is its output:</p>
<pre><code class="lang-plaintext">2.3.1
jac = ...
tensor([[[[ 0.1516],
          [ 0.0000],
          [ 0.0000]],

         [[-0.0771],
          [ 0.0000],
          [ 0.0000]]],


        [[[ 0.0000],
          [ 0.0703],
          [ 0.0000]],

         [[ 0.0000],
          [-0.0394],
          [ 0.0000]]],


        [[[ 0.0000],
          [ 0.0000],
          [ 0.1426]],

         [[ 0.0000],
          [ 0.0000],
          [-0.0732]]]])
grad = ...
tensor([[ 0.1516, -0.0771],
        [ 0.0703, -0.0394],
        [ 0.1426, -0.0732]])
gradB = ...
tensor([[ 0.1516, -0.0771],
        [ 0.0703, -0.0394],
        [ 0.1426, -0.0732]], grad_fn=&lt;SqueezeBackward0&gt;)
torch.equal (grad, gradB): True
</code></pre>
<p>*)  I used <code>jacfwd()</code> rather than <code>jacrev()</code> because you have fewer inputs<br>
(per batch element), namely one, than outputs, namely two (although with<br>
only two output elements, this won’t matter much).</p>
<p>Best.</p>
<p>K. Frank</p>