Grad returns repeating values for independent batch entries and doesn't want to second order derivative

username: JosueCom
cooked: <p>So, I have this function:</p>
<pre><code class="lang-python">  def ode_residual(self,
                 time: th.DoubleTensor,
                 state: th.DoubleTensor) -&gt; th.DoubleTensor:

      dtheta = th.autograd.grad(state.sum(), time, retain_graph=True)[0]
      d2theta = th.autograd.grad(dtheta.sum(), time, retain_graph=True)[0]
</code></pre>
<p>where <code>time</code> and 2 tensors are concatenated as an input to an MLP. <code>time</code> has requires_grad set to True and the other 2 variables do not. The shape of <code>time</code> is <code>(b, 1)</code> and the other 2 tensors have size <code>(b, 2)</code>. The MLP outputs a tensor that is sliced to 3 tensors. <code>state</code> is one of these tensors and has a shape of <code>(b, 1)</code> where b is the batch size.</p>
<pre><code class="lang-python">class MLP:

    def __init__(self,
                 width: int, 
                 depth: int,
                 pretrained:bool = False,
                 dtype=th.double) -&gt; None:
        super().__init__()

        self.layers = [nn.Linear(1+2+2, width, dtype=dtype)]
        self.layers.extend([nn.Linear(width, width, dtype=dtype) for _ in range(depth)])
        self.layers.append(nn.Linear(width, 1+ 1+ 2, dtype=dtype))

        self.layers = nn.Sequential(*self.layers)
    
    def output_to_states(self, output) -&gt; Tuple[th.DoubleTensor, th.DoubleTensor, th.DoubleTensor]:

        """
        
        Returns:
            state
            tensor 2
            tensor 3
        """

        return output[:, 0:1], output[:, 1:2], output[:, 2:]

    def forward(self, 
                t: th.DoubleTensor, 
                x0: th.DoubleTensor, 
                x1: th.DoubleTensor) -&gt; th.DoubleTensor:

        x = th.cat([t, x0, x1], dim=1)
        return self.layers(x)
</code></pre>
<p>First, in the code above, <code>dtheta</code> returns in a tensor of shape <code>(b, 1)</code> that has all entries be the same value. Why is this? Shouldn’t all batch entries be independent and thus even thought we have a summation, the gradients will be different for each batch. Also, calculating for <code>dtheta2</code> results in the following error even though I set <code>retain_graph</code> to True. Any help in answering these question would be appreciated.</p>
<pre><code class="lang-bash">File "c:\path\to\project\core\dynamics.py", line 93, in ode_residual
    d2theta = th.autograd.grad(dtheta.sum(), time, retain_graph=True)[0]
  File "C:\path\to\project\.venv\lib\site-packages\torch\autograd\__init__.py", line 303, in grad
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
</code></pre>

username: KFrank
cooked: <p>Hi Josue!</p>
<aside class="quote no-group quote-modified" data-username="JosueCom" data-post="1" data-topic="189997" 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/josuecom/48/27043_2.png" class="avatar"> JosueCom:</div>
<blockquote>
<pre data-code-wrap="python"><code class="lang-python">class MLP:
...
        self.layers = [nn.Linear(1+2+2, width, dtype=dtype)]
        self.layers.extend([nn.Linear(width, width, dtype=dtype) for _ in range(depth)])
        self.layers.append(nn.Linear(width, 1+ 1+ 2, dtype=dtype))

        self.layers = nn.Sequential(*self.layers)
</code></pre>
<p><code>dtheta</code> returns in a tensor of shape <code>(b, 1)</code> that has all entries be the same value. Why is this?</p>
</blockquote>
</aside>
<p><code>MLP</code> is basically a <code>Sequential</code> that consists solely of <code>Linear</code> layers <em>with<br>
no intervening non-linearities.</em>  It therefore collapses, in effect, into a single<br>
<code>Linear (1+2+2, 1+2+2)</code> (with a bunch of redundant parameters) <em>regardless</em><br>
of the values of <code>width</code> and <code>depth</code>.</p>
<p>So the <code>time</code> component of each batch element gets multiplied by the<br>
corresponding components of the collapsed <code>Linear</code>, those terms are<br>
summed (and the <code>bias</code> and terms from the “other 2 variables” are added).<br>
The gradient is just this single multiplicative factor (from the collapsed<br>
<code>Linear</code>) and the other terms <em>added</em> on don’t affect the gradient.  Because<br>
there is no non-additive “mixing” (no intervening non-linear “activations”),<br>
the “other 2 variables” (and the various additive terms) don’t affect <code>time</code>’s<br>
gradient, nor does the value of <code>time</code> itself affect its gradient, so the <code>time</code><br>
gradient is the same for all of the batch elements, even when the batch<br>
elements themselves differ.</p>
<aside class="quote no-group">
<blockquote>
<p>calculating for <code>dtheta2</code> results in the following error even though I set <code>retain_graph</code> to True.</p>
<pre data-code-wrap="python"><code class="lang-python">      dtheta = th.autograd.grad(state.sum(), time, retain_graph=True)[0]
      d2theta = th.autograd.grad(dtheta.sum(), time, retain_graph=True)[0]
</code></pre>
</blockquote>
</aside>
<p>In the <code>dtheta</code> line you should be using <code>create_graph = True</code> rather than<br>
<code>retain_graph = True</code>.  This causes the first call to <code>autograd.grad()</code> to<br>
<em>create</em> the graph that the second call to <code>autograd.grad()</code> will then use to<br>
compute the second derivative.  (Note that <code>create_graph = True</code> automatically<br>
turns on <code>retain_graph = True</code> so that the newly-created graph will be preserved<br>
for use by the second call to <code>autograd.grad()</code></p>
<p>(Your posted code, in and of itself, does not need <code>retain_graph = True</code> in its<br>
second call to <code>autograd.grad()</code>.)</p>
<p>Best.</p>
<p>K. Frank</p>

username: JosueCom
cooked: <p>Thank you so much! I knew of the problem with just having linear terms but for some reason I didn’t see it haha.</p>
<p>I have other calls to <code>grad()</code> outside of what I showed.</p>