How do pytorch deal with the sparse jacobian matrix in jvp/vjp during autograd?

username: KagamineLenOffical
cooked: <p>I’m using pytorch to deal with least square problem, and there is a step which need to get the jacobian for y w.r.t. x.<br>
This will cause huge memory usage since pytorch need to save the sparse jacobian matrix in dense mode(and other memory usage during get jacobian) although the calculation is quite simple.<br>
Like this:</p>
<pre><code class="lang-auto">x_ = torch.ones(100000,requires_grad=True)
def func(x__):
    return x__*2
y=func(x_)
print(torch.autograd.functional.jacobian(func,x_))
print(jacfwd(func)(x_))
</code></pre>
<p>And get error report:<br>
<code>DefaultCPUAllocator: not enough memory: you tried to allocate 40000000000 bytes.</code><br>
But we can still get the result of <code>autograd.grad()</code> and don’t use much memory. I’m curious about how pytorch deal with the jacobian matrix in jvp/vjp operation.<br>
In chumpy, the auto-diff just build the jacobian in sparse way. Does pytorch use similar method, or specify the way to process jvp for different operation?<br>
It’s quite hard to get memory usage under control since I don’t know how things work here.</p>

username: valerian.rey
cooked: <p>Hi, I’m also very interested in computing sparse Jacobians, and I have the same problem as you in 2025 with torch 2.7.0.</p>
<p>I tried the following ways of computing the Jacobian:</p>
<pre data-code-wrap="python"><code class="lang-python">jacobian = torch.autograd.functional.jacobian(func, x_, strategy="reverse-mode", vectorize=True)
jacobian = torch.autograd.functional.jacobian(func, x_, strategy="reverse-mode", vectorize=False)
jacobian = torch.autograd.functional.jacobian(func, x_, strategy="forward-mode", vectorize=True)
jacobian = torch.func.jacrev(func)(x_)
jacobian = torch.func.jacfwd(func)(x_)
jacobian = torch._functorch.deprecated.jacfwd(func)(x_)
</code></pre>
<p>And all of them lead to a huge memory usage (meaning that they don’t handle sparsity automatically).</p>
<blockquote>
<p>But we can still get the result of <code>autograd.grad()</code> and don’t use much memory. I’m curious about how pytorch deal with the jacobian matrix in jvp/vjp operation.</p>
</blockquote>
<p>I’m also curious about this. The internal vector-jacobian product that <code>torch.autograd.grad</code> has to make seems to be done without materializing the jacobian (of size 100k x 100k) in memory.</p>
<p>I could, however, obtain the Jacobian (this time even bigger: 10M x 10M) quite efficiently by providing a sparse identity matrix to a vmapped vjp:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
from torch import vmap
from torch.func import vjp

m = 10_000_000
x_ = torch.ones(m, requires_grad=True)

def func(x__):
  return x__ * 2

indices = torch.arange(0, m).expand((2, m))
values = torch.full((m,), 1.)
dentity_matrix = torch.sparse_coo_tensor(indices, values)
jacobian = vmap(vjp(func, x_)[1])(identity_matrix)

print(jacobian)
</code></pre>
<p>Note that the obtained jacobian matrix also has a <code>torch.sparse_coo</code> layout (same as <code>identity_matrix</code>).</p>
<p><strong>As a follow-up question, does anyone know if there is a more standard way to compute sparse jacobians? What I’m doing feels like a hack, and it does not seem to work for all sparse layouts. Any hint would be appreciated <img src="https://discuss.pytorch.org/images/emoji/apple/heart_hands.png?v=14" title=":heart_hands:" class="emoji" alt=":heart_hands:" loading="lazy" width="20" height="20"></strong></p>