Doubt regarding shape after Jacobian

vainaijr: <p>when I apply Jacobian in PyTorch autograd (torch==1.5.0), it gives me three dimensional output.<br>
for example,</p>
<pre><code class="lang-auto">def exp_reducer(x):
    return (x**2).sum(dim=1)
inputs = torch.rand(2, 2)
torch.autograd.functional.jacobian(exp_reducer, inputs)
</code></pre>
<pre><code class="lang-auto">tensor([[[0.3296, 0.9791],
         [0.0000, 0.0000]],

        [[0.0000, 0.0000],
         [1.6706, 0.7365]]])
</code></pre>
<p>and</p>
<pre><code class="lang-auto">torch.autograd.functional.jacobian(exp_reducer, inputs).shape
</code></pre>
<pre><code class="lang-auto">(2,2,2)
</code></pre>
<p>but I get a 2 dimensional output when using sympy,</p>
<pre><code class="lang-auto">x0, y0, x1, y1= sympy.symbols('x0, y0, x1, y1')
f = Matrix([Function('g')(x0, x1), Function('h')(y0, y1)])
k = Matrix([[x0, x1, y0, y1]])
f.jacobian(k)
</code></pre>
<p>then it give me, this,<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/6/d6bbf608961f07401cf5f6e12d44728e11f5cc8f.png" data-download-href="https://discuss.pytorch.org/uploads/default/d6bbf608961f07401cf5f6e12d44728e11f5cc8f" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/6/d6bbf608961f07401cf5f6e12d44728e11f5cc8f_2_690x113.png" alt="image" data-base62-sha1="uDCSLOTyafgnGxk4Xp5YSmxPQHt" width="690" height="113" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/6/d6bbf608961f07401cf5f6e12d44728e11f5cc8f_2_690x113.png, https://discuss.pytorch.org/uploads/default/optimized/3X/d/6/d6bbf608961f07401cf5f6e12d44728e11f5cc8f_2_1035x169.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/d/6/d6bbf608961f07401cf5f6e12d44728e11f5cc8f_2_1380x226.png 2x" data-dominant-color="353535"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1782×293 18.2 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>should the shape after applying Jacobian be 2d or 3d?</p>

albanD: <p>Hi,</p>
<p>Yes we follow a slightly different semantic.<br>
The usual definition for multivariate functions is given a 1D input and 1D output, you end up with as 2D Jacobian.</p>
<p>What sympy is doing here is to linearize the inputs from a (2, 2) Tensor into a (4,) Tensor. And so the Jacobian is (2, 4).</p>
<p>Because we usually allow multiple Tensors, people very often multiple dimensions and linearizing the Tensors can be confusing if you don’t know how they are layed out in memory, we made a different choice here (similar to what tensorflow or jax are doing).<br>
Here we don’t linearize Tensors and so you can index your Jacobian the same way you would index your original Tensor.</p>

vainaijr: <p>Would it be a good idea to add the symbolic representation to docs?<br>
When computing hessian, I get same dimension confusion,<br>
for example,</p>
<pre><code class="lang-auto">def pow_adder_reducer(x, y):
  print('x: ', x, 'y: ', y)
  return (2 * x.pow(2) + 3 * y.pow(2)).sum()
inputs = (torch.rand(2), torch.rand(2))
torch.autograd.functional.hessian(pow_adder_reducer, inputs)
</code></pre>
<p>PyTorch gives,</p>
<pre><code class="lang-auto">x:  tensor([0.4406, 0.9028], grad_fn=&lt;ViewBackward&gt;) y:  tensor([0.8790, 0.7052], grad_fn=&lt;ViewBackward&gt;)
((tensor([[4., 0.],
          [0., 4.]]), tensor([[0., 0.],
          [0., 0.]])), (tensor([[0., 0.],
          [0., 0.]]), tensor([[6., 0.],
          [0., 6.]])))
</code></pre>
<p>from sympy I get,<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/8/d854cb488227b539967694184824b7894a2d2d97.png" data-download-href="https://discuss.pytorch.org/uploads/default/d854cb488227b539967694184824b7894a2d2d97" title="Screenshot (531)"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/8/d854cb488227b539967694184824b7894a2d2d97_2_689x276.png" alt="Screenshot (531)" data-base62-sha1="uRKNQU02tcTEILQR8pjQIDP4TLF" width="689" height="276" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/8/d854cb488227b539967694184824b7894a2d2d97_2_689x276.png, https://discuss.pytorch.org/uploads/default/optimized/3X/d/8/d854cb488227b539967694184824b7894a2d2d97_2_1033x414.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/d/8/d854cb488227b539967694184824b7894a2d2d97.png 2x" data-dominant-color="323232"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot (531)</span><span class="informations">1346×539 19.3 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>is it doing something like,</p>
<pre><code class="lang-auto">f1 = 2*(x0**2) + 2*(x1**2)
f2 = 3*(y0**3) + 3*(y1**3)
</code></pre>
<pre><code class="lang-auto">hessian(f1, (x0, x1))
</code></pre>
<pre><code class="lang-auto">hessian(f1, (y0, y1))
</code></pre>
<pre><code class="lang-auto">hessian(f2, (x0, x1))
</code></pre>
<pre><code class="lang-auto">hessian(f2, (y0, y1))
</code></pre>
<p>that results in four 2x2 matrices.</p>

albanD: <p>Yes, this is a similar argument. The indexing is interleaved as output tensor, input tensor, output dim, input dim.<br>
You can check the discussion we had about this on the original issue here: <a href="https://github.com/pytorch/pytorch/issues/30632#issuecomment-575315939" class="inline-onebox">Add high level autograd functions · Issue #30632 · pytorch/pytorch · GitHub</a></p>
<blockquote>
<p>Would it be a good idea to add the symbolic representation to docs?</p>
</blockquote>
<p>Definitely open to suggestions to improve the doc !<br>
What do you mean here by “symbolic representation”? Could you give an example?</p>

vainaijr: <p>Like this one,</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/5/a/5a958736bac83f92a0b9ea247725aab4bc3bd05b.png" data-download-href="https://discuss.pytorch.org/uploads/default/5a958736bac83f92a0b9ea247725aab4bc3bd05b" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/5/a/5a958736bac83f92a0b9ea247725aab4bc3bd05b_2_690x112.png" alt="image" data-base62-sha1="cVlkLGztrA0G1zvdabqEq6SvEHF" width="690" height="112" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/5/a/5a958736bac83f92a0b9ea247725aab4bc3bd05b_2_690x112.png, https://discuss.pytorch.org/uploads/default/optimized/3X/5/a/5a958736bac83f92a0b9ea247725aab4bc3bd05b_2_1035x168.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/5/a/5a958736bac83f92a0b9ea247725aab4bc3bd05b.png 2x" data-dominant-color="353535"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1035×169 11.3 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>and this one,<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/b/ab97d00502e6b7ea068ccaae0cec22cd1473ed6b.png" data-download-href="https://discuss.pytorch.org/uploads/default/ab97d00502e6b7ea068ccaae0cec22cd1473ed6b" title="Screenshot (533)"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/a/b/ab97d00502e6b7ea068ccaae0cec22cd1473ed6b_2_690x124.png" alt="Screenshot (533)" data-base62-sha1="otYU4tT6p9tPMOlzoGSwHwjpGaL" width="690" height="124" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/a/b/ab97d00502e6b7ea068ccaae0cec22cd1473ed6b_2_690x124.png, https://discuss.pytorch.org/uploads/default/optimized/3X/a/b/ab97d00502e6b7ea068ccaae0cec22cd1473ed6b_2_1035x186.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/a/b/ab97d00502e6b7ea068ccaae0cec22cd1473ed6b_2_1380x248.png 2x" data-dominant-color="353535"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot (533)</span><span class="informations">1695×307 19.3 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>like this symbolic representation added to docs, this makes it a bit easier to understand when we input &gt;1d tensor.</p>

albanD: <p>Thanks, I added it to the issue to work on doc improvements here: <a href="https://github.com/pytorch/pytorch/issues/34998">https://github.com/pytorch/pytorch/issues/34998</a><br>
Feel free to post there if you have more ideas !</p>