How to calculate a jacobian for an entire batch

Yedidya_kfir: <p>I have a network which takes a vector size 10 and returns a vector size 20.<br>
I want to calculate the jacobian for the output of the network.<br>
I used the code</p>
<p><code>torch.autograd.functional.jacobian(func=network, inputs=x)</code><br>
to calculate it and it worked, I get the correct matrix size 20 * 10<br>
however when I try to do it over an entire batch (lets say, size 40) I get way over what I wanted (a matrix size 40 * 20 * 40 * 10</p>
<p>I believe the jacobian calculate the gradient of every batch w.r.t. to every input (ignoring the separation between batches)</p>
<p>How can I get the function to treat those batches as separate without using a for loop and separating them myself?</p>
<p>I think I can extract the jacobian by batches like so</p>
<pre><code class="lang-auto">indicies = torch.arange(x.shape[0])
jacobian = jacobian[indicies, :, indicies, :]
</code></pre>
<p>Is this correct? is there more efficient way to prevent pytorch calculating unnecessary derivitives?</p>

KFrank: <p>Hi Yedidya!</p>
<aside class="quote no-group" data-username="Yedidya_kfir" data-post="1" data-topic="198343" 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/yedidya_kfir/48/68245_2.png" class="avatar"> Yedidya_kfir:</div>
<blockquote>
<p>I believe the jacobian calculate the gradient of every batch w.r.t. to every input (ignoring the separation between batches)</p>
<p>How can I get the function to treat those batches as separate without using a for loop and separating them myself?</p>
</blockquote>
</aside>
<p>If I understand what you want to do, I believe that <a class="mention" href="/u/alphabetagamma96">@AlphaBetaGamma96</a>’s<br>
post that uses <code>vmap()</code> and <code>jacrev()</code> will give you your “batch jacobian:”</p>
<aside class="quote quote-modified" data-post="12" data-topic="80771">
  <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/a/3da27b/48.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/computing-batch-jacobian-efficiently/80771/12">Computing batch Jacobian efficiently</a> 
  </div>
  <blockquote>
    Just an update for anyone who reads the thread in the future, as of PyTorch2,the functorch library is now included in pytorch. So you can replace functorch with torch.func, for the most part the syntax is the same except if you have an nn.Module you’ll need to create a ‘functional’ version of your model. 
For example, 
model = myModel(*args, **kwargs) #our network

from torch.func import vmap, jacrev, functional_call

params = dict(model.named_parameters())
inputs = torch.randn(batch_size, input…
  </blockquote>
</aside>

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

Yedidya_kfir: <aside class="quote no-group" data-username="KFrank" data-post="2" data-topic="198343">
<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"> KFrank:</div>
<blockquote>
<p>erstand what you want</p>
</blockquote>
</aside>
<p>This solution indeed seems far more efficient than mine. however I must admit I didnt get how it works.<br>
what is the difference between the jacobian function and the jacrev function? also how the vmap function works?<br>
Thanks anyway, you cut my performance in 70%</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/yedidya_kfir">@Yedidya_kfir</a>, and <a class="mention" href="/u/kfrank">@KFrank</a>!</p>
<p>The way that the <code>torch.func</code> namespace works is by creating a so-called functionalized version of your model that is now a <a href="https://en.wikipedia.org/wiki/Pure_function" rel="noopener nofollow ugc">pure function</a> of only inputs and outputs (i.e. no other non-local variables). This is why your network parameters are now an input (and not an internal variable like before).</p>
<p>The <code>jacrev</code> function is <code>torch.func</code>’s version of <code>torch.autograd.functional.jacobian</code> and allows for the computation of reverse-mode jacobian calculations. This will allow you compute the jacobian of a functionalized module within the <code>torch.func</code> namespace.</p>
<p>Now, why did I suggest to use the <code>torch.func</code> namespace and not just <code>torch.autograd.functional.jacobian</code>?</p>
<p>Well, we can compute all samples in parallel  with <code>torch.func.vmap</code>. What this function does is it redefines your function to work for a <strong>single</strong> sample, and then uses this transformation to compute all samples in parallel. This is where the speed-up occurs, instead of calling the autograd engine <code>N</code> times (for  <code>N</code>  samples) via a for-loop, you effectively only call the autograd engine <strong>once</strong> in a single vectorized call and compute all samples in parallel.</p>
<p>There are limitations to the <code>torch.func</code> namespace, so I’ve left a link to the docs <a href="https://pytorch.org/docs/stable/func.ux_limitations.html#vmap-limitations" rel="noopener nofollow ugc">here</a>.</p>