Optimize objective involving jacobian

af_c: <p>Suppose i have two function z = f(x), g(z), and i want to<br>
z1 = f(x1)<br>
z2 = f(x2)<br>
optimize MSE(\gradient(g) (z1), \gradient(g) (z2)).<br>
But when i backward propogate the gradient of parameter of f is all zero.</p>
<p>Here is the code</p>
<pre><code class="lang-auto">import torch
from torch import nn
def batch_jacobian(func, x, create_graph=True, strict=True):
    def _func_sum(x):
        return func(x).sum(dim=0)
    return torch.autograd.functional.jacobian(_func_sum, x, create_graph=create_graph, strict=strict).permute(1,0,2)

g = nn.Sequential(nn.Linear(2, 4), nn.ReLU(), nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 3))
x1 = torch.rand((2, 3))
x2 = torch.rand((2, 3))
mseloss = nn.MSELoss()

theta = torch.arange(6).reshape(3, 2) / 10
theta.requires_grad = True
print(theta)
sig = torch.nn.Sigmoid()
def f(x):
    return sig(x @ theta)

opt = torch.optim.Adam([theta] + list(g.parameters()), lr=1e-3)
opt.zero_grad()
loss = mseloss(batch_jacobian(g, f(x1)), batch_jacobian(g, f(x2)))
loss.backward()
opt.step()
</code></pre>
<p>And the output of grad of theta is all zero</p>

KFrank: <p>Hi af!</p>
<aside class="quote no-group quote-modified" data-username="af_c" data-post="1" data-topic="219202" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/af_c/48/76690_2.png" class="avatar"> af_c:</div>
<blockquote>
<p>But when i backward propogate the gradient of parameter of f is all zero.<br>
…</p>
<pre><code class="lang-auto">def batch_jacobian(func, x, create_graph=True, strict=True):
    def _func_sum(x):
        return func(x).sum(dim=0)
    return torch.autograd.functional.jacobian(_func_sum, x, create_graph=create_graph, strict=strict).permute(1,0,2)
...
theta = torch.arange(6).reshape(3, 2) / 10
theta.requires_grad = True
...
def f(x):
    return sig(x @ theta)
...
loss = mseloss(batch_jacobian(g, f(x1)), batch_jacobian(g, f(x2)))
</code></pre>
<p>And the output of grad of theta is all zero</p>
</blockquote>
</aside>
<p>The problem is that <code>torch.autograd.functional.jacobian()</code> only backpropagates back to<br>
its <code>inputs</code> argument (your <code>f (x1)</code> and <code>f (x2)</code>).  It neither knows nor cares that, say, <code>f (x1)</code><br>
depends on <code>theta</code>, you don’t backpropagate through <code>f (x1)</code>, and therefore you never reach<br>
the dependence on <code>theta</code>, so you get no <code>.grad</code> for <code>theta</code>.</p>
<p>I’ve tweaked your script so that the call to <code>f (x)</code> occurs inside of <code>batch_jacobian()</code> and<br>
therefore inside of the call to <code>torch.autograd.functional.jacobian()</code>.  Doing so does then<br>
produce <code>.grad</code> for <code>theta</code>.</p>
<p>Here is the tweaked script:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

torch.manual_seed (2025)

from torch import nn

theta = torch.arange(6).reshape(3, 2) / 10
theta.requires_grad = True
print ('theta = ...')
print (theta)
sig = torch.nn.Sigmoid()
def f(x):
    return sig(x @ theta)

def batch_jacobian(func, x, create_graph=True, strict=True):
    def _func_sum(x):
        # return func(x).sum(dim=0)
        return func (f (x)).sum (dim=0)
    return torch.autograd.functional.jacobian(_func_sum, x, create_graph=create_graph, strict=strict).permute(1,0,2)

g = nn.Sequential(nn.Linear(2, 4), nn.ReLU(), nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 3))
x1 = torch.rand((2, 3))
x2 = torch.rand((2, 3))
mseloss = nn.MSELoss()

opt = torch.optim.Adam([theta] + list(g.parameters()), lr=1e-3)
opt.zero_grad()
# loss = mseloss(batch_jacobian(g, f(x1)), batch_jacobian(g, f(x2)))
loss = mseloss(batch_jacobian(g, x1), batch_jacobian(g, x2))
loss.backward()
opt.step()

print ('loss:', loss)
print ('theta = ...')
print (theta)
print ('theta.grad = ...')
print (theta.grad)
</code></pre>
<p>And here is its output:</p>
<pre><code class="lang-plaintext">2.6.0+cu126
theta = ...
tensor([[0.0000, 0.1000],
        [0.2000, 0.3000],
        [0.4000, 0.5000]], requires_grad=True)
loss: tensor(3.6073e-08, grad_fn=&lt;MseLossBackward0&gt;)
theta = ...
tensor([[0.0010, 0.0990],
        [0.2009, 0.2990],
        [0.4009, 0.4990]], requires_grad=True)
theta.grad = ...
tensor([[-2.0444e-07,  4.2059e-07],
        [-1.0772e-07,  2.4320e-07],
        [-1.5494e-07,  3.4528e-07]])
</code></pre>
<p>Best.</p>
<p>K. Frank</p>