Autograd and Temporary Variables

username: Vaker
cooked: <p>I’m fairly new to PyTorch, and I’m running into some questions about how autograd will interact with some parts of my code which I haven’t been able to find answers to.</p>
<pre><code class="lang-auto">  def forward(self,theta: Tensor) -&gt; Tensor: 
    #theta is a (B,F,L) tensor, where B is a batch dimension
    th0 = theta[:,:,0].view(self.b,5,1,1,1)
    th1 = theta[:,:,1].view(self.b,5,1,1,1)
    th2 = theta[:,:,2].view(self.b,5,1,1,1)
    th3 = theta[:,:,3].view(self.b,5,1,1,1)
    
    #frames is a 5-tensor of shape (B,F,L,4,4)
    #in the start state, all of the (4,4) blocks are identity
    frames = self.start_state.clone()
    
    #swivel is a 5-tensor of shape (B,F,L,4,4)
    S = self.swivel.clone()
    
    S[...,:3,:3] = self.I + torch.sin(th0)*S[...,:3,:3] + \
      (1-torch.cos(th0))*(S[...,:3,:3]@S[...,:3,:3])
      
    S.requires_grad = True
    S.retain_grad()
    
    frames = S@frames
</code></pre>
<p>I had a few questions about this. After some similar further transformations, the frames tensor (or a particular slice of it) will be the output.</p>
<ol>
<li>
<p>S here is a temporary derived from a member variable of my class- from some reading I got the impression that <code>S.retain_grad()</code> would allow the gradients of S to be maintained for the backward pass, but is this correct? If I returned <code>frames</code> at this step and called a backward pass, it ‘works’, but I’m not sure if it’s actually computing gradients through <code>S</code> as I would like it to.</p>
</li>
<li>
<p>If I put the two lines enabling grad for S before the operation on the 3x3 slice of its last two dimensions, PyTorch gives an error saying an in-place operation was called on a leaf of the compute graph. This makes sense, and the current placement raises no errors, but I’m left wondering if putting those lines <em>after</em> that op is actually fixing the issue, or if it will still interact weirdly with the autograd engine?</p>
</li>
<li>
<p>I’m not very sure about how editing/multiplying slices of tensors affects gradient computation, if at all. For example, the next transformations to be applied to <code>frames</code> are:</p>
</li>
</ol>
<pre><code class="lang-auto"># curl is a (B,F,L-1,4,4) tensor
C1 = self.curl.clone() 
  ....
    
C2 = self.curl[:,:,:2,...].clone()
  ....
frames[:,:,1:,:,:] = C1@frames[:,:,1:,:,:]
frames[:,:,2:,:,:] = C2@frames[:,:,2:,:,:]

#I've cut out a few lines computing blocks of C1, C2 and enabling grad here, 
#they're almost identical to what was done for S
</code></pre>
<p>My intent is to only apply the transformations of C1, C2 on sections of <code>frames</code> in a differentiable manner- but will this work, or is there a nuance I’m missing?</p>
<p>Any help would be greatly appreciated! I’ve not worked with PyTorch internals much at all (although I’m slowly growing through the tutorials and docs), so any detail might help.</p>
<p>(P.S. - if you have stylistic comments, I’d be glad to hear them as well. I have tried to do things cleanly where possible, but I’m not sure that I’ve succeeded whatsoever.)</p>

username: KFrank
cooked: <p>Hi Vaker!</p>
<aside class="quote no-group quote-modified" data-username="Vaker" data-post="1" data-topic="219882" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/v/da6949/48.png" class="avatar"> Vaker:</div>
<blockquote>
<pre><code class="lang-auto">  def forward(self,theta: Tensor) -&gt; Tensor: 
    #theta is a (B,F,L) tensor, where B is a batch dimension
    th0 = theta[:,:,0].view(self.b,5,1,1,1)
    ...
    frames = self.start_state.clone()
    
    #swivel is a 5-tensor of shape (B,F,L,4,4)
    S = self.swivel.clone()
    
    S[...,:3,:3] = self.I + torch.sin(th0)*S[...,:3,:3] + \
    ...
    S.requires_grad = True
    S.retain_grad()
</code></pre>
</blockquote>
</aside>
<p>It’s not clear what you are trying to do or what is going on here.</p>
<p>Some specific questions:</p>
<p>Is autograd supposed to be tracking gradients for <code>theta</code>?  Specifically, does the <code>theta</code><br>
passed into <code>forward()</code> carry <code>requires_grad = True</code>, as would be the case if it were the<br>
output of some previous layer with trainable parameters?</p>
<p>Does <code>self.start_state</code> carry <code>requires_grad = True</code>?</p>
<p>Does <code>swivel</code> carry <code>requires_grad = True</code>?  (Presumably not.)</p>
<aside class="quote no-group">
<blockquote>
<ol>
<li>S here is a temporary derived from a member variable of my class- from some reading I got the impression that <code>S.retain_grad()</code> would allow the gradients of S to be maintained for the backward pass, but is this correct?</li>
</ol>
</blockquote>
</aside>
<p>If I understand what you are asking, not exactly.  <code>S.requires_grad = True</code> would tell<br>
autograd to track gradients for <code>S</code> (while <code>S.retain_grad()</code> – which you would typically<br>
not use – does something rather different).</p>
<aside class="quote no-group">
<blockquote>
<p>If I returned <code>frames</code> at this step and called a backward pass, it ‘works’, but I’m not sure if it’s actually computing gradients through <code>S</code> as I would like it to.</p>
</blockquote>
</aside>
<p>Autograd will not compute gradients “through” <code>S</code>.  It will only compute gradients back up<br>
to (and including) <code>S</code>, but not beyond.  This is because, as your code stands, <code>S</code> is a<br>
so-called “leaf variable” – the point at which autograd <em>begins</em> tracking gradients during<br>
the forward pass.</p>
<aside class="quote no-group">
<blockquote>
<ol start="3">
<li>If I put the two lines enabling grad for S before the operation on the 3x3 slice of its last two dimensions, PyTorch gives an error saying an in-place operation was called on a leaf of the compute graph.</li>
</ol>
</blockquote>
</aside>
<p>I deduce (but without the relevant code cannot verify) that prior to the statement<br>
<code>S.requires_grad = True</code>, <code>S.requires_grad</code> has the value <code>False</code>.  Therefore, <code>S</code> is<br>
not yet a leaf variable of the computation graph, so modifying it inplace by executing<br>
the line <code>S[...,:3,:3] = ...</code> is allowed.  But when you move <code>S.requires_grad = True</code><br>
to before <code>S[...,:3,:3] = ...</code>, <code>S</code> now is already a leaf variable prior to your attempted<br>
inplace modification, so you get the error.</p>
<aside class="quote no-group">
<blockquote>
<p>I’m left wondering if putting those lines <em>after</em> that op is actually fixing the issue</p>
</blockquote>
</aside>
<p>It won’t fix the issue that I imagine you are thinking about (but I don’t really know what<br>
you are trying to do).  Concretely, gradients from <code>frames = S @ frames</code> will not flow<br>
back up <em>through</em> <code>S</code> to <code>self.swivel</code> because (as written) <code>S</code> is the leaf variable, so<br>
that’s where the backward pass will stop.</p>
<p>(As an aside, if <code>self.start_state</code> carries <code>requires_grad = True</code>, gradients <em>will</em><br>
flow back up through the <code>frames</code> branch of the computation graph.  The fact that <code>S</code><br>
is a leaf variable only blocks backpropagation back up to variables from which it was<br>
computed – not other parts of the computation graph, such as the frames branch.)</p>
<p>Best.</p>
<p>K. Frank</p>

username: Vaker
cooked: <p>Hi Frank, thanks for your reply!</p>
<p>I apologize for the lack of clarity in my question- I tried to post a minimal code example to illustrate the types of operations I was confused about, but in retrospect I can see how I’ve left rather too much for the reader to infer.</p>
<p>For some context, this function is intended to implement batched forward kinematics for a particular class of parallel robotic manipulators. The <code>theta</code> tensor can be interpreted as a batch of joint (angle) configurations. In my application, the joint angles are predicted by a neural network, which I’ll call <code>net</code>. Everything stored in <code>self</code> can be viewed as an expression of the structure of the manipulator; so they’re all constant. I’ll give more detail below.</p>
<p>(Now answering your specific questions)</p>
<ol>
<li>
<p>You’re correct about <code>theta</code>- it’s an output from <code>net</code>, and the output of the <code>forward</code> function here is intended for use in a loss for training <code>net</code>- so it has grad enabled.</p>
</li>
<li>
<p><code>self.start_state</code> is a collection of constant transforms defining a neutral position of the robot; it has <code>requires_grad = False</code>.</p>
</li>
<li>
<p><code>self.swivel</code> is just an expansion of a single constant 4x4 matrix- if that matrix is <code>A</code>, then<br>
<code>self.swivel = A.view(1,1,1,4,4).expand(B,F,L,-1,-1)</code>, where <code>B, F, L</code> are the same dimensions as for <code>theta</code> in my original post.<br>
In my code it’s used to generate the (<code>theta</code>-dependent) transform <code>S</code>. So <code>swivel</code> has <code>requires_grad = False</code> because it’s a constant, and I’m assuming that the multiplication by elements of <code>theta</code> makes S differentiable- but I’m not sure if that’s true.<br>
Working in the simple case B = F = L = 1,  <code>S = I + sin(theta)*A + (1-cos(theta))*A@A</code>, and I’d ideally want <code>S' = cos(theta)*A + sin(theta)*A@A</code></p>
</li>
</ol>
<p>Replying to the rest of your message:</p>
<blockquote>
<p>If I understand what you are asking, not exactly. <code>S.requires_grad = True</code> would tell<br>
autograd to track gradients for <code>S</code> (while <code>S.retain_grad()</code> – which you would typically<br>
not use – does something rather different).</p>
</blockquote>
<p>Ah, alright- I read up on <code>.retain_grad()</code> again, and after checking <code>S.is_leaf()</code> I can see that it’s not necessary. I had thought that temporary variables needed to be explicitly added as leaves of the computation graph, but that was probably not giving enough credit to the pytorch folks.</p>
<blockquote>
<p>Autograd will not compute gradients “through” <code>S</code>. It will only compute gradients back up<br>
to (and including) <code>S</code>, but not beyond. This is because, as your code stands, <code>S</code> is a<br>
so-called “leaf variable” – the point at which autograd <em>begins</em> tracking gradients during<br>
the forward pass.</p>
</blockquote>
<p>Right, I think that make sense. As I expanded above, a slice of the output of this <code>forward</code> (essentially just <code>frames</code> matrix multiplied by a sequence of transformations, the first of which is <code>S</code>) will be put through a simple, differentiable <code>loss</code> function- at <em>that</em> point, when I call <code>loss.backward()</code>, would it be correct to say that autograd will compute gradients through S? Or am I still misunderstanding?</p>
<aside class="quote no-group">
<blockquote>
<p>I deduce (but without the relevant code cannot verify) that prior to the statement<br>
<code>S.requires_grad = True</code>, <code>S.requires_grad</code> has the value <code>False</code>. Therefore, <code>S</code> is<br>
not yet a leaf variable of the computation graph, so modifying it inplace by executing<br>
the line <code>S[...,:3,:3] = ...</code> is allowed. But when you move <code>S.requires_grad = True</code><br>
to before <code>S[...,:3,:3] = ...</code>, <code>S</code> now is already a leaf variable prior to your attempted<br>
inplace modification, so you get the error.</p>
</blockquote>
</aside>
<p>Yes, your deductions are absolutely right, and this explanation makes a lot of sense.</p>
<blockquote>
<p>It won’t fix the issue that I imagine you are thinking about (but I don’t really know what<br>
you are trying to do). Concretely, gradients from <code>frames = S @ frames</code> will not flow<br>
back up <em>through</em> <code>S</code> to <code>self.swivel</code> because (as written) <code>S</code> is the leaf variable, so<br>
that’s where the backward pass will stop.</p>
</blockquote>
<p>With the new context that <code>self.swivel</code> is constant, am I right to say that this part does not matter? I only need gradients to flow back through <code>theta</code> for the training of <code>net</code>- will that actually work? If my statement that <code>S</code> is no longer a leaf after the loss computation is correct, then I think it will? (I apologize again here- I cannot find a way to state this more precisely.)</p>
<aside class="quote no-group">
<blockquote>
<p>(As an aside, if <code>self.start_state</code> carries <code>requires_grad = True</code>, gradients <em>will</em><br>
flow back up through the <code>frames</code> branch of the computation graph. The fact that <code>S</code><br>
is a leaf variable only blocks backpropagation back up to variables from which it was<br>
computed – not other parts of the computation graph, such as the frames branch.)</p>
</blockquote>
</aside>
<p>Huh, that’s not quite how I thought it worked. I had thought that since the <code>self.start_state</code> is constant, <code>self.start_state.requires_grad = False</code> would just treat it as a constant in gradient calculation, rather than not having gradients flow through it at all.</p>
<p>I think all in all I need to do loads more reading into how autograd actually works, alongside some practice.  My past projects have been simple enough where this sort of thing hasn’t been an issue, but I can’t see myself doing anything particularly sophisticated with so much missing knowledge.</p>
<p>Thanks again for your response- it was really very helpful, and it’s given me a good number of avenues for learning. If you have any further comments, I’d be glad to hear them!</p>
<p>Best regards,<br>
Vaker</p>

username: KFrank
cooked: <p>Hi Vaker!</p>
<p>Let me make a couple of comments:</p>
<p>First, you should play around with some simple computations (e.g., <code>w = x + y * z</code>)<br>
and track how things like <code>requires_grad = True</code> evolve through the intermediate<br>
results, depending on which of the initial variables start out with <code>requires_grad = True</code>.</p>
<p>There is a package – that I have never used – called <a href="https://pypi.org/project/torchviz/" rel="noopener nofollow ugc">torchviz</a> that will supposedly help<br>
you visualize computation graphs.  You might play around with this to get a feel for what<br>
computation graphs look like in practice.</p>
<p>Let me add some clarification to my mention of <code>.is_leaf</code>:  I’m not sure how <code>.is_leaf</code><br>
is used in pytorch – to me, it means “has the potential to become the leaf of an actual<br>
computation graph.”  I say this because <code>.is_leaf</code> seems to be <code>True</code> for tensors that<br>
have <code>requires_grad = False</code>, even if they haven’t been used in any computation.</p>
<p>Consider:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'2.7.0+cu128'
&gt;&gt;&gt; t1 = torch.ones (3)
&gt;&gt;&gt; s1 = t1 * t1
&gt;&gt;&gt; t1.requires_grad
False
&gt;&gt;&gt; t1.is_leaf
True
&gt;&gt;&gt; s1.requires_grad
False
&gt;&gt;&gt; s1.is_leaf
True
&gt;&gt;&gt; t2 = torch.ones (3, requires_grad = True)
&gt;&gt;&gt; s2 = t2 * t2
&gt;&gt;&gt; t2.requires_grad
True
&gt;&gt;&gt; t2.is_leaf
True
&gt;&gt;&gt; s2.requires_grad
True
&gt;&gt;&gt; s2.is_leaf
False
</code></pre>
<p>That is, a tensor that has <code>requires_graph = False</code> isn’t (in my mind) a full-fledged<br>
member – <em>through which gradients flow</em> – of any actual computation graph (leaf or not),<br>
but it has <code>.is_leaf = True</code>.  (Note, in <code>c = 1.2</code>, <code>y = c * x</code>, we could call <code>c</code> a leaf<br>
of the computation graph because it is needed to compute the gradient of <code>y</code> with respect<br>
to <code>x</code>.  But <code>c</code>, itself, doesn’t have any <code>.grad</code> and no gradients flow <em>through</em> it.)  Just<br>
something to bear in mind if you’re using <code>.is_leaf</code> to track computation graphs.</p>
<aside class="quote no-group" data-username="Vaker" data-post="3" data-topic="219882" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/v/da6949/48.png" class="avatar"> Vaker:</div>
<blockquote>
<ol>
<li>
<p>You’re correct about <code>theta</code>- it’s an output from <code>net</code>, and the output of the <code>forward</code> function here is intended for use in a loss for training <code>net</code>- so it has grad enabled.</p>
</li>
<li>
<p><code>self.start_state</code> is a collection of constant transforms defining a neutral position of the robot; it has <code>requires_grad = False</code>.</p>
</li>
<li>
<p><code>self.swivel</code> is just an expansion of a single constant 4x4 matrix</p>
</li>
</ol>
</blockquote>
</aside>
<p>Going back to the code you posted initially:</p>
<aside class="quote no-group">
<blockquote>
<pre data-code-wrap="python"><code class="lang-python">  def forward(self,theta: Tensor) -&gt; Tensor: 
    th0 = theta[:,:,0].view(self.b,5,1,1,1)
    ...
    frames = self.start_state.clone()
    S = self.swivel.clone()
</code></pre>
</blockquote>
</aside>
<p>As I understand you, the <code>theta</code> passed in has <code>requires_grad = True</code>, while <code>start_state</code><br>
and <code>swivel</code> do not.  So at this point, neither <code>frames</code> nor <code>S</code> has <code>requires_grad = True</code>.</p>
<aside class="quote no-group">
<blockquote>
<pre data-code-wrap="python"><code class="lang-python">    S[...,:3,:3] = self.I + torch.sin(th0)*S[...,:3,:3] + \
      (1-torch.cos(th0))*(S[...,:3,:3]@S[...,:3,:3])
</code></pre>
</blockquote>
</aside>
<p>However, <code>th0</code> does have <code>requires_grad = True</code>, so after the above-quoted line, <code>S</code><br>
will, as well.</p>
<aside class="quote no-group">
<blockquote>
<pre data-code-wrap="python"><code class="lang-python">    S.requires_grad = True
    S.retain_grad()
    
    frames = S@frames
</code></pre>
</blockquote>
</aside>
<p>As I read your code, <code>S.requires_grad = True</code> does nothing (as <code>.requires_grad</code> was<br>
already <code>True</code>) (and <code>S.retain_grad()</code> probably doesn’t do anything you want).</p>
<p>Finally, <code>frames = S@frames</code> creates a new tensor (that happens to use the same<br>
python reference name, “frames,” that used to refer to the “old” <code>frames</code>).  This new<br>
<code>frames</code> will have <code>requires_grad = True</code> because <code>S</code> did (even though the old<br>
<code>frames</code> didn’t).</p>
<aside class="quote no-group">
<blockquote>
<p>I’m assuming that the multiplication by elements of <code>theta</code> makes S differentiable</p>
</blockquote>
</aside>
<p>Yes, this is true.  <code>S</code> should be “differentiable” in that it carries <code>requires_grad = True</code>.</p>
<aside class="quote no-group">
<blockquote>
<p>I only need gradients to flow back through <code>theta</code> for the training of <code>net</code>- will that actually work?</p>
</blockquote>
</aside>
<p>Let me correct what I said in my earlier post.  Based on what you’ve said, after<br>
<code>S[...,:3,:3] = self.I + torch.sin(th0)*S[...,:3,:3] + ...</code>, <code>S</code> will<br>
have <code>requires_grad = True</code> and depend on <code>th0</code> and thus on <code>theta</code>.  So if<br>
you compute your loss value using the (new) <code>frames</code>, calling <code>.backward()</code> on<br>
that loss value will backpropagate gradients back up through <code>S</code> to <code>theta</code> (as I<br>
understand is what you want).</p>
<p>Given that <code>S</code> already has <code>requires_grad = True</code>, <code>S.requires_grad = True</code><br>
is a no-op, and you can (and stylistically should) leave it out.</p>
<aside class="quote no-group">
<blockquote>
<p>If my statement that <code>S</code> is no longer a leaf after the loss computation is correct</p>
</blockquote>
</aside>
<p>In my usage, <code>S</code> was never a leaf variable (in the sense of being a leaf of some<br>
actual computation graph).  Immediately after <code>S = self.swivel.clone()</code>, <code>S</code> does<br>
have <code>.is_leaf = True</code>, but because is also has <code>requires_grad = False</code>, I don’t<br>
count is as a “real” leaf variable.</p>
<p>Then after <code> S[...,:3,:3] = self.I + torch.sin(th0)*S[...,:3,:3] + ...</code> <code>S</code><br>
becomes a non-leaf interior node of the computation graph (through which gradients<br>
will flow).</p>
<p>Consider:</p>
<pre data-code-wrap="python"><code class="lang-python">&gt;&gt;&gt; t = torch.ones (3, requires_grad = True)
&gt;&gt;&gt; s = torch.ones (3)
&gt;&gt;&gt; s.requires_grad
False
&gt;&gt;&gt; s[1] = t[2]
&gt;&gt;&gt; s.requires_grad
True
&gt;&gt;&gt; s.is_leaf
False
</code></pre>
<p>Note, computing a loss value from (the new) <code>frames</code> doesn’t change any of this.  <code>S</code><br>
is still an interior node of the graph – computing the loss just adds a new “root” to<br>
the “head” of the graph.</p>
<p>Best.</p>
<p>K. Frank</p>

username: Vaker
cooked: <p>Hi Frank! Sorry for the delayed reply, had an unusually packed weekend.</p>
<p>i) I’ll take a look at torchviz- if it works as advertised, it seems as though it could be a good tool to understand the construction of torch computational graphs in situations where I’m unsure of how it will work.<br>
ii) The point on <code>.is_leaf</code> is surprising- I can certainly see myself using it incorrectly and getting confused if you hadn’t pointed it out.</p>
<p>To the rest of your reply, your comments on the position of <code>S</code> in the graph structure make a lot of sense, and taking those alongside the examples you provided and my own reading I don’t think I have any follow up questions.  (I’ve also verified on my end that everything works as you stated, and updated my code to remove non-operations.)</p>
<p>Thank you again for your detailed responses and explanations to my questions! They cleared up a good number of autograd pain points for me which might otherwise have taken a lot more work to understand.</p>
<p>(This response feels rather short, in contrast; I want to emphasize that this is because I really have no further questions, and I feel that this exchange has given me a solid starting point to learn more from on my own.)</p>