How to calculate Jacobians for a batch

username: windows95
cooked: <p>Hello, I want to calculate Jacobian matrices for a batch of data.</p>
<p>I have <code>x</code> (batch_size, 3) and calculated <code>y</code> (batch_size, 3). I need a Jacobian matrix of shape (batch_size, 3, 3)</p>
<p>I tried the following code:</p>
<pre><code class="lang-auto">import torch

x = torch.randn(1024, 3, requires_grad=True) # a batch of coordinates
y = torch.sin(x) # calculated y in forward pass

grads = []
for i_dim in range(y.shape[1]):
    y_i = y[:, i_dim]
    ones = torch.ones_like(y_i)
    # `torch.autograd.grad` seems to implicitly add dimensions up.
    # So I separately calculate grad for every dim of `y`
    # But this still fails.
    grad = torch.autograd.grad(y_i, x, grad_outputs=ones, create_graph=True, retain_graph=True, is_grads_batched=True)[0]
    grads.append(grad)

grad = torch.cat(grad, dim=1)

print(grad.shape)
</code></pre>
<pre><code class="lang-auto">RuntimeError: If `is_grads_batched=True`, 
we interpret the first dimension of each grad_output as the batch dimension. 
The sizes of the remaining dimensions are expected to match the shape of corresponding output, 
but a mismatch was detected: grad_output[0] has a shape of torch.Size([]) and output[0] has a shape of torch.Size([1024]). 
If you only want some tensors in `grad_output` to be considered batched, consider using vmap.
</code></pre>
<p>I think <code>outputs</code> and <code>grad_output</code> have already the same shape. Why is this error raised?</p>
<p>BTW, I tried <code>torch.func.jacrev()</code>, but I need to calculate <code>y</code> beforehand, not in the closure of <code>jacrev()</code>. My code has been written this way.</p>

username: KFrank
cooked: <p>Hi Windows!</p>
<aside class="quote no-group quote-modified" data-username="windows95" data-post="1" data-topic="217805">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/w/a8b319/48.png" class="avatar"> windows95:</div>
<blockquote>
<p>I need a Jacobian matrix of shape (batch_size, 3, 3)</p>
<pre><code class="lang-auto">RuntimeError: If `is_grads_batched=True`, 
we interpret the first dimension of each grad_output as the batch dimension. 
</code></pre>
<p>I think <code>outputs</code> and <code>grad_output</code> have already the same shape. Why is this error raised?</p>
</blockquote>
</aside>
<p>I don’t think that <code>grad_output</code> works quite the way you are expecting.</p>
<p>Here is a script, loosely based on the code you posted, that I think does what you want<br>
and also illustrates how you can use <code>vmap</code> / <code>jacrev</code> to compute a batch of jacobians:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

_ = torch.manual_seed (2025)

x = torch.randn (1024, 3, requires_grad = True)   # a batch of coordinates

# we use the fact that y[j] is independent of x[i] unless j == i
# (but we don't use that y[j, l] is independent of x[i, k] unless also l == k)
y = torch.sin (x)                                 # calculated y in forward pass

ysum = y.sum (dim = 0)                            # use independence of y[j] on x[i]

# is_grad_batched is batching over ysum's single dimension of length 3
grad_outputs = torch.eye (y.shape[1])
batch_grad = torch.autograd.grad (ysum, x, grad_outputs = grad_outputs, is_grads_batched = True)[0]

batch_grad = batch_grad.transpose (0, 1)          # move batch_size = 1024 to the front

print ('batch_grad.shape =', batch_grad.shape)
m = 77
print ('m =', m, ', batch_grad[m] = ...')         # observe that any batch_grad[m] is diagonal
print (batch_grad[m])


# use vmap and jacrev to compute jacobian (as batch_grad_b)

batch_jacobian = torch.vmap (torch.func.jacrev (torch.sin))
batch_grad_b = batch_jacobian (x)

print ('batch_grad_b.shape =', batch_grad_b.shape)
print ('torch.equal (batch_grad, batch_grad_b) =', torch.equal (batch_grad, batch_grad_b)) 
</code></pre>
<p>And here is its output:</p>
<pre><code class="lang-plaintext">2.6.0+cu126
batch_grad.shape = torch.Size([1024, 3, 3])
m = 77 , batch_grad[m] = ...
tensor([[0.4986, 0.0000, 0.0000],
        [0.0000, 0.9996, 0.0000],
        [0.0000, 0.0000, 0.8100]])
batch_grad_b.shape = torch.Size([1024, 3, 3])
torch.equal (batch_grad, batch_grad_b) = True
</code></pre>
<p>It’s not clear to me which of the two approaches illustrated will run faster.  The first<br>
approach uses a single explicit forward pass and, under the hood, three backward<br>
passes that I think are <code>vmap</code>ed.  The second approach, however, runs 1024 <code>vmap</code>ed<br>
forward passes, so you might expect that it would be significantly slower.</p>
<aside class="quote no-group">
<blockquote>
<p>BTW, I tried <code>torch.func.jacrev()</code>, but I need to calculate <code>y</code> beforehand, not in the closure of <code>jacrev()</code>.</p>
</blockquote>
</aside>
<p>You don’t explain why you need to calculate <code>y</code> beforehand, but pytorch provides<br>
<a href="https://pytorch.org/docs/2.6/generated/torch.func.functional_call.html" rel="noopener nofollow ugc">functional_call()</a> that makes it easy to run a model as a closure.</p>
<p>Best.</p>
<p>K. Frank</p>

username: windows95
cooked: <p>Thank you very much. Now I understand when samples in a batch are independent, summing <code>y</code> over batch dim is OK.</p>
<p>Then another thought comes to me: what if the samples in a batch are not independent? I have used attention in <code>model</code>, so the sample coordinates are not independent. Can I still compute a batch of Jacobians like <code>(batch_size, y_dim, x_dim)</code>, not <code>(batch_size, y_dim, batch_size, x_dim)</code>, using <code>torch.autograd.grad()</code>?</p>
<p>I’m not sure how to use <code>functional_call()</code> yet, but I can explain the problem I encounter when trying <code>jacrev()</code>.</p>
<p>In real code it is a PINN problem, what I send to <code>model()</code> is not only a tensor <code>x</code>, but a tuple of a bunch of things, like another tensor and multiple masks. And <code>model</code>’s output is also not a single tensor.</p>
<p>My code is like <code>out1, out2 = model(x1, x2, mask1, mask2)</code> and then compute Jacobian of <code>out1</code> with respect to <code>x1</code>. After that <code>out1</code> and <code>out2</code> are still needed for other computation.</p>

username: KFrank
cooked: <p>Hi Windows!</p>
<aside class="quote no-group" data-username="windows95" data-post="3" data-topic="217805">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/w/a8b319/48.png" class="avatar"> windows95:</div>
<blockquote>
<p>the sample coordinates are not independent. Can I still compute a batch of Jacobians like <code>(batch_size, y_dim, x_dim)</code>, not <code>(batch_size, y_dim, batch_size, x_dim)</code>, using <code>torch.autograd.grad()</code>?</p>
</blockquote>
</aside>
<p>Yes, but because the “samples aren’t independent,” you can’t sum over the batch<br>
dimension of your output <code>y</code>.  You will have to have <code>autograd.grad()</code> perform 3072<br>
backward passes using <code>is_grads_batched = True</code>.  Also, technically speaking, you<br>
will not be keeping full jacobian matrices, but rather selecting out just the desired<br>
submatrices of the jacobian matrices (which will be computed in full).</p>
<p>Here is an example:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

_ = torch.manual_seed (2025)

fc = torch.nn.Linear (1024 * 3, 1024 * 3)          # connect all inputs to all outputs

x = torch.randn (1024, 3, requires_grad = True)    # a batch of coordinates

y = fc (x.flatten()).sigmoid().reshape (1024, 3)   # all elements depend on all elements

grad_outputs = torch.eye (1024 * 3).reshape (1024 * 3, 1024, 3)
batch_grad = torch.autograd.grad (y, x, grad_outputs = grad_outputs, is_grads_batched = True)[0]
print ('batch_grad.shape =', batch_grad.shape)

# select out the 3x3 "jacobian" matrices I think you want
batch_grad = batch_grad.reshape (1024, 3, 1024, 3).diagonal (dim1 = 0, dim2 = 2).permute (2, 0, 1)
print ('batch_grad.shape =', batch_grad.shape)
</code></pre>
<p>And here is its output:</p>
<pre><code class="lang-plaintext">2.6.0+cu126
batch_grad.shape = torch.Size([3072, 1024, 3])
batch_grad.shape = torch.Size([1024, 3, 3])
</code></pre>
<aside class="quote no-group">
<blockquote>
<p>I’m not sure how to use <code>functional_call()</code> yet, but I can explain the problem I encounter when trying <code>jacrev()</code>.<br>
…<br>
My code is like <code>out1, out2 = model(x1, x2, mask1, mask2)</code> and then compute Jacobian of <code>out1</code> with respect to <code>x1</code>. After that <code>out1</code> and <code>out2</code> are still needed for other computation.</p>
</blockquote>
</aside>
<p>You can package your call to model as a function (that refers to some “global” variables)<br>
and pass it to <code>jacrev()</code>:</p>
<pre data-code-wrap="python"><code class="lang-python"># x1 = whatever ...
# x2 = whatever ...
# mask1 = whatever ...
# mask2 = whatever ...

def applyModelToX1 (x_one):
    out1, _ = model (x_one, x2, mask1, mask2)
    return out1

fullJacobian = torch.func.javrev (applyModelToX1) (x1)

# snip out submatrices of fullJacobian, as desired
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

username: windows95
cooked: <p>Thanks Frank. Your code is right. I tested it in my actual implementation. Unfortunately, in my real code, <code>batch_size</code> is significantly large (about 10^6). Adding this additional dimension makes the real code unable to work.</p>
<p>I don’t know if there’s a way I can modify my model output into two versions. One is original output for standard backpropagation. The other, whose cross-sample derivative dependencies are eliminated or ignored, is going to be used only for Jacobian computation.</p>
<p>Although I’m not sure how to process the second output.</p>
<p>Is it a viable idea?</p>

username: KFrank
cooked: <p>Hi Windows!</p>
<aside class="quote no-group" data-username="windows95" data-post="5" data-topic="217805">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/w/a8b319/48.png" class="avatar"> windows95:</div>
<blockquote>
<p>in my real code, <code>batch_size</code> is significantly large (about 10^6). Adding this additional dimension makes the real code unable to work.</p>
</blockquote>
</aside>
<p>You don’t say how your real code doesn’t “work.”  Does it give the wrong answer?  Does<br>
it crash?  Do you run into memory problems?  Does it take too long?</p>
<p>As an aside, it’s hard to give useful answers when you keep moving the goal posts.  It<br>
appears that your real use case is a far cry from the example you gave initially:</p>
<pre data-code-wrap="python"><code class="lang-python">x = torch.randn (1024, 3, requires_grad = True)
y = torch.sin (x)
</code></pre>
<p>Does every element of <code>out1</code>, say <code>out1[0, 0]</code>, depend on more or less every element<br>
of <code>x1</code>?  Furthermore, do the derivatives of <code>out1[0, 0]</code> depend on more or less  every<br>
element of <code>x1</code>?</p>
<p>Can you perform a single complete forward pass – with your real <code>batch_size</code> – with<br>
autograd  turned off, e.g:</p>
<pre data-code-wrap="python"><code class="lang-python">with torch.inference_mode():
    out1, out2 = model (x1, x2, mask1, mask2)
</code></pre>
<p>Can you perform a full forward pass with autograd left on?</p>
<p>Can you perform a single backward pass (for a single scalar):</p>
<pre data-code-wrap="python"><code class="lang-python">out1, out2 = model (x1, x2, mask1, mask2)
loss = out1[0, 0]
loss.backward()
</code></pre>
<p>You could then just keep the piece of the jacobian you need,<br>
<code>j = x1.grad[0].detach.clone()</code>, and then loop over all the elements of <code>out1</code>,<br>
performing a backward pass for each element.  Before doing this, you could set<br>
<code>requires_grad = False</code> for <code>x2</code>, <code>mask1</code>, <code>mask2</code>, and all of the parameters of<br>
<code>model</code> to reduce the work and storage needed by the forward pass and backward<br>
passes.</p>
<p>If you can get this to work, you could probably get some speed-up with judicious use<br>
of <code>vmap()</code> and / or <code>autograd.grad (..., is_grads_batched = True)</code> so that your<br>
python loop (if you have one at all) doesn’t have to loop over each element of <code>out1</code><br>
individually.</p>
<p>Good luck.</p>
<p>K. Frank</p>