Is there a faster way to compute the jacobian than autograd.functional.jacobian?

Stev1: <p>Given a big model, a vector-valued function f:R^n → (f_1, f_2,…,f_m) with n&gt;&gt;m and m&gt;1. Now I need grad(f_i) w.r.t. to ONE SINGLE input x and I need it for every f_i. I can do that very easily just by using <em>autograd.functional.jacobian</em></p>
<pre><code class="lang-auto">jacobianMatrix = functional.jacobian(myFunction, x)
</code></pre>
<p>And from here on I can access all the needed gradients for f_i.</p>
<p>My question is, if there maybe is a faster method to do exactly that. I am aware of <em>torch.func</em> and <em>jacrev</em> and its vectorization capabilities but in my case where I just use one single input instance, <em>jacrev</em> seems a little bit slower and seemingly has a big memory footprint.</p>
<p>On the other hand I could use <em>autograd.grad</em> and just compute the jacobian^T vector product and use the unitvectors as <em>grad_outputs</em> to extract the gradients from the jacobian. I tested this version too and it is slower than just to compute the full jacobian.</p>
<p><strong>Is there maybe another method I could try?</strong></p>
<p>Why can I only compute gradients w.r.t. one instance? I don’t know the other inputs in advance, because input1 is calculated from the gradients w.r.t. input0 and therefore vectorization won’t help me in this case.</p>

KFrank: <p>Hi Stev!</p>
<aside class="quote no-group" data-username="Stev1" data-post="1" data-topic="218647">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/7ab992/48.png" class="avatar"> Stev1:</div>
<blockquote>
<p>Now I need grad(f_i) w.r.t. to ONE SINGLE input x and I need it for every f_i.</p>
</blockquote>
</aside>
<p>It’s not clear to me what you’re asking.</p>
<p>If by “ONE SINGLE input x” you mean that you have a function from R^n to R^m,<br>
f (x_1, x_2, …, x_n) → (f_1 (x_1, x_2, …, x_n), …, f_m (x_1, …, x_2, …, x_n)), and<br>
you want the partial derivatives with respect to just one of the x_i, say, x_7, then you<br>
should use <a href="https://pytorch.org/tutorials/intermediate/forward_ad_usage.html" rel="noopener nofollow ugc">forward-mode automatic differention.</a>  If you create a “dual tensor” for x_7,<br>
you will be able to compute all of the derivatives d f_i (x_j) / d x_7 using just a single<br>
forward-mode forward pass.</p>
<p>(If this is not what you are asking, please clarify.)</p>
<p>Best.</p>
<p>K. Frank</p>

Stev1: <p>Hey Frank!</p>
<p>I just need d f_i(x)/dx (for all i) where x is a vector. With the “one single input” I wanted to emphasize that I don’t look for the fastest method to compute this quantity for batches but the fastest for just one input.</p>
<p>Best.</p>
<p>Stev</p>

KFrank: <p>Hi Stev!</p>
<aside class="quote no-group" data-username="Stev1" data-post="3" data-topic="218647">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/7ab992/48.png" class="avatar"> Stev1:</div>
<blockquote>
<p>I just need d f_i(x)/dx (for all i) where x is a vector.</p>
</blockquote>
</aside>
<p>Okay, I think I understand now.</p>
<p>You said that the dimensionality of your input, n, is much greater than the dimensionality<br>
of your output, m.  Therefore you want to use regular “reverse-mode” autograd, rather<br>
than forward-mode autograd.</p>
<p><code>autograd.functional.jacobian()</code> (with <code>vectorize = True</code> and the default value for<br>
<code>strategy</code> of <code>'reverse-mode'</code>) should be the fastest approach.  Note that the documentation<br>
suggests that <a href="https://pytorch.org/docs/2.6/generated/torch.func.jacrev.html" rel="noopener nofollow ugc">jacrev()</a> is “less experimental,” but the two should be doing basically the same<br>
thing under the hood, namely using a loop to perform a separate backward pass for each<br>
element of the output, but with that loop “vectorized” using <code>vmap()</code>.</p>
<p>Best.</p>
<p>K. Frank</p>

Stev1: <p>Hi Frank!</p>
<aside class="quote no-group" data-username="KFrank" data-post="4" data-topic="218647">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/k/ecb155/48.png" class="avatar"> KFrank:</div>
<blockquote>
<p>Note that the documentation<br>
suggests that <a href="https://pytorch.org/docs/2.6/generated/torch.func.jacrev.html" rel="noopener nofollow ugc">jacrev()</a> is “less experimental,” but the two should be doing basically the same<br>
thing under the hood,</p>
</blockquote>
</aside>
<p>In my case <em>jacrev()</em> really makes my notebook laggy and I had several kernel crashes running it. I think that I have read somewhere that <em>jacrev</em> in contrast to <em>autograd.functional.jacobian()</em> sacrifices more memory for faster execution?</p>
<p>Otherwise, thank you for your help!</p>
<p>Best.</p>
<p>Stev</p>