The behavior of torch.autograd.functional.jvp

username: hhxx
cooked: <p>I am reading the documentation about <a href="https://pytorch.org/docs/stable/generated/torch.autograd.functional.jvp.html" rel="noopener nofollow ugc">jvp</a>. I believe it is used to compute the dot product between the Jacobian and v. However, I find its behavior strange when the inputs are 2-D matrices.</p>
<p>According to its usage: <code>torch.autograd.functional.jvp(func, inputs, v=None)</code>, we can define the func and inputs as follows:</p>
<pre><code class="lang-auto">def func(x): 
    return 2*x.sum(dim=1)
x = torch.ones(3,2, requiers_grad=True)
v = torch.tensor([[1,2],[3,4],[5,6])
res = torch.autograd.functional.jvp(func, x, v)
</code></pre>
<p>After running the above code:</p>
<pre><code class="lang-auto">res[0] = tensor([4,4,4]) # which is the output of func(x)
res[1] = tensor([6, 14, 22]) # which is the dot product between jacobian and v
</code></pre>
<p>The result of <strong>res[1]</strong> is somewhat unintuitive. Let’s define <strong>y = func(x)</strong>, which means <strong>y = 2*x.sum(dim=1)</strong> :</p>
<pre><code class="lang-auto">y[0] = 2 * x[0,0] + 2 * x[0,1]
y[1] = 2 * x[1,0] + 2 * x[1,1]
y[2] = 2 * x[2,0] + 2 * x[2,1]
</code></pre>
<p>The Jacobian matrix of y <em>with respect to</em> x is:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/7/e7458a013781dec94c267ea3eda4b080ea71fd59.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/e7458a013781dec94c267ea3eda4b080ea71fd59" title="1"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/e/7/e7458a013781dec94c267ea3eda4b080ea71fd59_2_517x195.jpeg" alt="1" data-base62-sha1="wZVhvNZ5Rp2yNyHZPsW1qtUu3up" width="517" height="195" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/e/7/e7458a013781dec94c267ea3eda4b080ea71fd59_2_517x195.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/e/7/e7458a013781dec94c267ea3eda4b080ea71fd59_2_775x292.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/e/7/e7458a013781dec94c267ea3eda4b080ea71fd59_2_1034x390.jpeg 2x" data-dominant-color="EAEAEA"><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">1</span><span class="informations">1071×406 128 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>
<pre><code class="lang-auto">grad = torch.tenosr([(2,2),(0,0),(0,0),
                     (0,0),(2,2),(0,0),
                     (0,0),(0,0),(2,2)])
grad.shape = torch.Size([3,3,2])
v = torch.tensor([[1,2],
                  [3,4],
                  [5,6]])
</code></pre>
<p>According to res[1] = tensor([6, 14, 22]), we can infer that the <strong>dot product</strong> is performed in the following way:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/4/e481f2858e21a3947dec832bf8b3cc68ea86d251.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/e481f2858e21a3947dec832bf8b3cc68ea86d251" title="1"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/e/4/e481f2858e21a3947dec832bf8b3cc68ea86d251_2_690x124.jpeg" alt="1" data-base62-sha1="wBtgZVXfs6WVh5keKdsa9gEQ28V" width="690" height="124" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/e/4/e481f2858e21a3947dec832bf8b3cc68ea86d251_2_690x124.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/e/4/e481f2858e21a3947dec832bf8b3cc68ea86d251_2_1035x186.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/e/4/e481f2858e21a3947dec832bf8b3cc68ea86d251.jpeg 2x" data-dominant-color="EAE8E8"><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">1</span><span class="informations">1266×228 113 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>Such an operation is not the standard <strong>dot product</strong> between two matrices. Additionally, the output of JVP only has the shape of a torch.Size([3]), and the inputs of JVP have the shape of a torch.Size([3, 3]). So I wonder what the meaning of the outputs of JVP is? What does the grad of res[1] stand for, and where can we use it?</p>

username: hhxx
cooked: <p>I have realized the motivation behind why <code>jvp</code> performs the Jacobian-vector-product operation with 1-D array inputs: It is designed to compute the partial gradients of two variables using the chain rule. In such cases, <code>res[1]</code> also has the same shape as the inputs. It is clear that <code>res[1]</code> represents the gradients of the inputs. However, this doesn’t hold when the inputs are a 2-D matrix, which greatly confused me.</p>

username: albanD
cooked: <p>Hey!</p>
<p>Indeed, this definition of computing a dot product only makes sense for 1D inputs.<br>
But you can also view your 2D input as a 1D input (.view(-1)), compute the jacobian vector product and then reshape again the result to be the shape of the output.<br>
In this view, the jacobian has size nb_elements_outputs x nb_elements_inputs which shows the linearization.</p>