Calculating the Jacobian matrix after calculating the gradient with autograd

username: SadeghS
cooked: <p>Hi everyone,</p>
<p>I’ve been trying to calculate the jacobian matrix or Jacobian times a vector when the explicit formula for the gradient is not available, and I calculate it by autograd. In this case, the jacobian of PyTorch returns a zero matrix. In the matrix-vector product, it returns the same value as the gradient of the objective function with respect to parameter theta. Please find a toy example below. I appreciate your guidance.</p>
<p><span class="hashtag">#First</span> version<br>
<span class="hashtag">#Forming</span> the Jacobian matrix<br>
def objective(x,theta):<br>
return torch.exp(theta[0])*x[0]*x[1] +torch.exp(theta[1])*x[2]</p>
<p>def gradient(x,theta):<br>
funval = objective(x,theta)<br>
funval.backward()<br>
return x.grad</p>
<p>x = torch.tensor([1.0,2.0,3.0], requires_grad=True)<br>
theta = torch.tensor([-1.0,-3.0], requires_grad=True)<br>
jac = jacobian(gradient,(x,theta))<br>
print(jac)</p>
<p><span class="hashtag">#Second</span> version<br>
<span class="hashtag">#Forming</span> the Jacobian matrix-vector product with, for example, the vector of ones<br>
def objective(x,theta):<br>
return torch.exp(theta[0])*x[0]*x[1] +torch.exp(theta[1])*x[2]</p>
<p>def gradient(x,theta):<br>
funval = objective(x,theta)<br>
funval.backward()<br>
print(theta.grad)<br>
return x.grad</p>
<p>x = torch.tensor([1.0,2.0,3.0], requires_grad=True)<br>
theta = torch.tensor([-1.0,-3.0], requires_grad=True)<br>
grad = Variable(gradient(x,theta),requires_grad=True)<br>
temp = torch.dot(grad,torch.ones(3))<br>
temp.backward()<br>
print(theta.grad)</p>
<p><span class="hashtag">#which</span> is the same as the gradient of the objective function with respect to theta</p>

username: KFrank
cooked: <p>Hi Sadegh!</p>
<aside class="quote no-group quote-modified" data-username="SadeghS" data-post="1" data-topic="160665" 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/sadeghs/48/52933_2.png" class="avatar"> SadeghS:</div>
<blockquote>
<p>I’ve been trying to calculate the jacobian matrix<br>
…<br>
def objective(x,theta):<br>
return torch.exp(theta[0])*x[0]*x[1] +torch.exp(theta[1])*x[2]</p>
</blockquote>
</aside>
<p>The <em>Jacobian</em> of a <strong>vector-valued</strong> function of a vector argument is<br>
the matrix of partial derivatives of each element of the function’s value<br>
with respect to each element of the argument.</p>
<p>In your case your <code>objective()</code> function returns a scalar value, not a<br>
vector, so we normally would not use the term Jacobian.  (If you want<br>
to say that your scalar result is a one-dimensional vector, you could<br>
treat your gradient vector as a 1 x n matrix and call it the Jacobian.)</p>
<p>Could you clarify what you are asking here?</p>
<p>Note that the <em>Hessian</em> is the  matrix of second-order mixed partial<br>
derivatives of a scalar-valued function of a vector argument.  Could<br>
that be what you are asking about?</p>
<p>Best.</p>
<p>K. Frank</p>

username: SadeghS
cooked: <p>Thank you for your reply, K. Frank.</p>
<p>Let me clarify. Let f(x,theta) be our objective. At first, I calculate its gradient with respect to x when theta is constant in the function gradient(x,theta) which is a vector-valued function. Then, I want to calculate the derivative of the function gradient(x,theta) with respect to theta. The result will be the Jacobian matrix. Did it make it clear?</p>

username: KFrank
cooked: <p>Hi Sadegh!</p>
<aside class="quote no-group" data-username="SadeghS" data-post="3" data-topic="160665" 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/sadeghs/48/52933_2.png" class="avatar"> SadeghS:</div>
<blockquote>
<p>Then, I want to calculate the derivative of the function gradient(x,theta) with respect to theta. The result will be the Jacobian matrix.</p>
</blockquote>
</aside>
<p>If you have a vector-valued function, <code>g()</code> (that happens to be the gradient<br>
of some scalar-valued function, <code>f()</code>), then yes, the derivative of <code>g()</code> will<br>
be <code>g()</code>'s Jacobian.</p>
<p>But that’s an odd way of describing is.  You are, of course, computing the<br>
second-order derivatives of <code>f()</code>, which is to say, you are computing the<br>
Hessian of <code>f()</code>.  Calling it the Jacobian (of <code>g()</code>) just obscures what it going<br>
on.</p>
<p>Regardless of what you choose to call it, the computation you describe<br>
is that of the mixed second-order partial derivatives of <code>f()</code>.</p>
<p>I would suggest that you start with pytorch’s <a href="https://pytorch.org/docs/1.12/generated/torch.autograd.functional.hessian.html" rel="noopener nofollow ugc">hessian()</a> functional to<br>
compute this.</p>
<p>It is true that <code>hessian()</code> will compute the full Hessian of <code>f()</code>, whereas<br>
in your description you only want the <code>x</code>-<code>theta</code> cross terms, so there could<br>
be some inefficiency in that you would compute the block-diagonal terms<br>
that you don’t want.</p>
<p>But I would recommend <code>hessian()</code>, only moving on to something more<br>
complicated if it proves inadequate for your needs.</p>
<p>Best.</p>
<p>K. Frank</p>