Shape of tensors to return for backward()

sus: <p>I’m trying to write a PyTorch Function for some black box function we’ll call f(x,y,z), where x,y,z are vectors of varying length and f returns a vector of length 4. I’m confused on the dimensions that I should be returning for the backward function. For example’s sake, we’ll say that x is a tensor with a single dimension of length 2, then in the backwards function I return a tensor with dimensions 2x4, since the vector that f returns is of length 4. Am I correct on this? Or is there some other way that I’m supposed to calculate the gradient that ALSO results in a tensor that is a single dimension of length 2, i.e. compress the output of [f(x+dx)-f(x)]/dx into a scalar?</p>

KFrank: <p>Hi Sus!</p>
<aside class="quote no-group" data-username="sus" data-post="1" data-topic="184786" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/ea5d25/48.png" class="avatar"> sus:</div>
<blockquote>
<p>I’m trying to write a PyTorch Function for some black box function we’ll call f(x,y,z), where x,y,z are vectors of varying length and f returns a vector of length 4. I’m confused on the dimensions that I should be returning for the backward function.</p>
</blockquote>
</aside>
<p>In short, your backward function should return a tuple of tensors that<br>
individually have the same shapes as the tensors input to your custom<br>
(forward) function.</p>
<p>Quoting from <a href="https://pytorch.org/docs/2.0/notes/extending.html#extending-torch-autograd" rel="noopener nofollow ugc">Extending <code>torch.autograd</code>:</a></p>
<blockquote>
<p><code>backward()</code> (or <code>vjp()</code>) defines the gradient formula.<br>
…<br>
It should return as many tensors as there were inputs, with each of them containing the gradient w.r.t. its corresponding input.</p>
</blockquote>
<p>Just to be explicit, the “gradient w.r.t. its corresponding input” will have<br>
the same shape and type as that input.</p>
<aside class="quote no-group">
<blockquote>
<p>x is a tensor with a single dimension of length 2, then in the backwards function I return a tensor with dimensions 2x4, since the vector that f returns is of length 4. Am I correct on this?</p>
</blockquote>
</aside>
<p>No.  In your example case, the backward function should return a<br>
one-dimensional tensor of length 2 (to match the input x, and two<br>
more gradient tensors that match y and z).</p>
<aside class="quote no-group">
<blockquote>
<p>Or is there some other way that I’m supposed to calculate the gradient that ALSO results in a tensor that is a single dimension of length 2</p>
</blockquote>
</aside>
<p>Yes.  Your backward function is supposed to compute the so-called<br>
<em>vector-Jacobian product.</em>  Your backward function will have passed into<br>
it a tensor (more precisely, a tuple of tensors) that is the same shape as<br>
the output of your custom function.  This is the “vector” in vector-Jacobian<br>
product.</p>
<p>In your example case, this will be a one-dimensional tensor of length 4 (to<br>
match the shape of the output of f).  The Jacobian of f (with respect to its<br>
first argument, x) will indeed be a tensor of shape 2x4, but your backward<br>
function is supposed to <em>contract</em> the 2x4 Jacobian (if it even explicitly<br>
computes the Jacobian, which it need not necessarily do) with the length-4<br>
vector passed into it to form the length-2 vector-Jacobian product that it then<br>
returns.</p>
<p>Best.</p>
<p>K. Frank</p>

sus: <p>Thank you so much for the detailed response! From what you’ve told me and what I’ve looked up online about the vector-Jacobian product, I would multiply each partial derivative (row of the 2x4?) of f by the output of f as the row vector, thus making the 2x4 into a 2x1 tensor, then I could just reshape it to be my vjp for x?</p>
<p>Also, what do you mean that I don’t need to explicitly compute the Jacobian? I think that’s exactly what I’m doing and it would be great if you knew about any resources that could clear that up. Thank you so much!</p>

KFrank: <p>Hi Sus!</p>
<aside class="quote no-group" data-username="sus" data-post="3" data-topic="184786" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/ea5d25/48.png" class="avatar"> sus:</div>
<blockquote>
<p>I would multiply each partial derivative (row of the 2x4?) of f by the output of f as the row vector, thus making the 2x4 into a 2x1 tensor</p>
</blockquote>
</aside>
<p>As you’ve stated things, no.</p>
<p>The <em>output</em> of f (that is, the output f produced during the forward<br>
pass) is not directly relevant.  Autograd keeps track of the entire<br>
computation graph and, in particular, of everything that happens<br>
to the output of f during the rest of the forward pass.</p>
<p>In the backward pass, autograd passes into f’s backward function<br>
a “vector” that represents the gradient of the final loss function with<br>
respect to the output of f.</p>
<p>The <a href="https://pytorch.org/docs/2.0/generated/torch.autograd.Function.backward.html" rel="noopener nofollow ugc">documentation for Function.backward()</a> calls this argument <code>grad_outputs</code>.</p>
<p>So, in your case, your Jacobian matrix would have shape <code>[2, 4]</code><br>
and <code>grad_outputs</code> would have shape <code>[4]</code>.  You would multiply<br>
<code>grad_outputs</code> by the Jacobian matrix to produce the return value<br>
of <code>backward()</code> of shape <code>[2]</code>.  (So I think that your result tensor will<br>
be a one-dimensional vector without any “singleton” dimension that<br>
you would need to <code>.reshape()</code> or <code>.squeeze()</code> away.)</p>
<p>(For simplicity, I’m only talking about the case where f outputs only<br>
a single tensor so that autograd only passes a single <code>grad_outputs</code><br>
tensor to your <code>backward()</code> function.)</p>
<aside class="quote no-group">
<blockquote>
<p>what do you mean that I don’t need to explicitly compute the Jacobian?</p>
</blockquote>
</aside>
<p>You’re certainly allowed to compute the Jacobian and multiply it onto<br>
<code>grad_outputs</code>.  But <code>backward()</code> only has to return the <em>result</em> of that<br>
multiplication.  If it can construct that result by some other (possibly<br>
cheaper) means, it’s not required to construct the Jacobian explicitly<br>
nor explicitly perform the multiplication.</p>
<p>As a trivial example, suppose your custom function maps a length-<code>n</code><br>
vector to another length-<code>n</code> vector by multiplying it by <code>12.0</code>.  The<br>
Jacobian of this function is an <code>n x n</code> diagonal matrix that is <code>12.0</code><br>
times the identity matrix. <code>backward()</code> can multiply <code>grad_outputs</code><br>
by the scalar <code>12.0</code> and return it – no need to “materialize” an <code>n x n</code><br>
matrix nor perform a matrix-vector multiplication.</p>
<p>Best.</p>
<p>K. Frank</p>

sus: <blockquote>
<p>The <em>output</em> of f (that is, the output f produced during the forward pass) is not directly relevant. Autograd keeps track of the entire computation graph and, in particular, of everything that happens to the output of f during the rest of the forward pass.</p>
</blockquote>
<p>Whoops! That is exactly what I thought I was typing, got tripped up by the <code>grad_outputs</code> variable name I have in my code. Thank you!</p>
<blockquote>
<p>If it can construct that result by some other (possibly cheaper) means, it’s not required to construct the Jacobian explicitly nor explicitly perform the multiplication.</p>
</blockquote>
<p>Oh, I see, thank you so much! It definitely clears things up more for me</p>