Autograd loses track of gradient if nested

username: panos
cooked: <p>I have the following function</p>
<p>template<br>
torch::Tensor ppppppH(const torch::Tensor &amp;x,<br>
const torch::Tensor &amp;p,<br>
T W,<br>
std::function&lt;torch::Tensor(const torch::Tensor&amp;, const torch::Tensor&amp;, T)&gt; H)<br>
{<br>
// Create tensors with gradient tracking for both x and p<br>
auto xt = x.clone();<br>
xt.detach_();<br>
xt=xt.set_requires_grad(false);<br>
auto pt = p.clone();<br>
pt.detach_();<br>
pt=pt.set_requires_grad(true);</p>
<pre><code>std::cerr &lt;&lt; "After setting gradient in ppppppH" &lt;&lt; std::endl;
std::cerr &lt;&lt; "xt.requires_grad() = " &lt;&lt; xt.requires_grad() &lt;&lt; std::endl;
std::cerr &lt;&lt; "pt.requires_grad() = " &lt;&lt; pt.requires_grad() &lt;&lt; std::endl;

// Compute the Hamiltonian value
auto Hvalue = H(xt, pt, W);
std::cerr &lt;&lt; "Hvalue.requires_grad() = " &lt;&lt; Hvalue.requires_grad() &lt;&lt; std::endl;

// Compute the first-order gradient of H with respect to p
auto grad_H_p = torch::autograd::grad({Hvalue}, 
                                      {pt}, 
                                      {torch::ones_like(Hvalue)}, 
                                      true, 
                                      true, 
                                      true)[0];
// Initialize a tensor to hold the third-order gradients
std::cerr &lt;&lt; "grad_H_p.requires_grad() = " &lt;&lt; grad_H_p.requires_grad() &lt;&lt; std::endl;
std::cerr &lt;&lt; "grad_H_p.grad_fn() = " &lt;&lt; grad_H_p.grad_fn() &lt;&lt; std::endl;
std::cerr &lt;&lt; "grad_H_p shape = " &lt;&lt; grad_H_p.sizes() &lt;&lt; std::endl;
std::cerr &lt;&lt; "grad_H_p values = " &lt;&lt; grad_H_p &lt;&lt; std::endl;

auto batch_size = p.size(0);
auto p_dim = p.size(1);
torch::Tensor third_order_derivative = torch::zeros({batch_size, p_dim, p_dim, p_dim}, p.options());

// Compute the second-order and third-order gradients
for (int i = 0; i &lt; p_dim; ++i)
{
    // Zero out the gradients of pt before computing the gradient for the next dimension
    if (pt.grad().defined()) {
        pt.grad().zero_();
    }

    // Compute the gradient of the i-th component of grad_H_p with respect to p
    auto grad_H_pati = grad_H_p.index({Slice(), i});
    auto grad_H_p_i = torch::autograd::grad({grad_H_pati}, 
                                            {pt}, 
                                            {torch::ones_like(grad_H_pati)}, 
                                            true, 
                                            true, 
                                            true)[0];
    std::cerr &lt;&lt; "grad_H_p_i.requires_grad() = " &lt;&lt; grad_H_p_i.requires_grad() &lt;&lt; std::endl;
    std::cerr &lt;&lt; "pt.requires_grad() = " &lt;&lt; pt.requires_grad() &lt;&lt; std::endl;
    for (int j = 0; j &lt; p_dim; ++j)
    {
        // Zero out the gradients of pt before computing the gradient for the next dimension
        if (pt.grad().defined()) {
            pt.grad().zero_();
        }
        //Check to see if grad_H_p_i lost the graph
        if ( grad_H_p_i.requires_grad() == false) {
            std::cerr &lt;&lt; "grad_H_p_i lost gradient tracking at this point." &lt;&lt; std::endl;
            grad_H_p_i.set_requires_grad(true);
        }


        // Compute the gradient of the j-th component of grad_H_p_i with respect to p
        auto grad_H_p_ij = torch::autograd::grad({grad_H_p_i.index({Slice(), j})}, 
                                                 {pt}, 
                                                 {torch::ones_like(grad_H_p_i.index({Slice(), j}))}, 
                                                 true, 
                                                 false, 
                                                 true)[0];
        std::cerr &lt;&lt; "grad_H_p_ij.requires_grad() = " &lt;&lt; grad_H_p_ij.requires_grad() &lt;&lt; std::endl;

        if (!grad_H_p_ij.requires_grad()) {
            std::cerr &lt;&lt; "grad_H_p_ij lost gradient tracking at this point." &lt;&lt; std::endl;
            continue;
        }

        //third_order_derivative.select(1, i).select(2, j).copy_(grad_H_p_ij);
        third_order_derivative.index_put_({Slice(), i, j}, grad_H_p_ij);
    }
}

// Return the third-order derivative tensor
return third_order_derivative;
</code></pre>
<p>}</p>
<p>The function is tested using the following methods<br>
torch::Tensor control(const torch::Tensor&amp; x1,<br>
const torch::Tensor&amp; x2,<br>
const torch::Tensor&amp; p1,<br>
const torch::Tensor&amp; p2,<br>
double W=1.0) {<br>
auto u = -p2 * ((1 - x1 * x1) * x2 - x1) / W;<br>
return u; // Return through copy elision<br>
}</p>
<p>torch::Tensor hamiltonian(const torch::Tensor&amp; x,<br>
const torch::Tensor&amp; p,<br>
double W) {<br>
torch::Tensor p1 = p.index({Slice(), 0});<br>
torch::Tensor p2 = p.index({Slice(), 1});<br>
torch::Tensor x1 = x.index({Slice(), 0});<br>
torch::Tensor x2 = x.index({Slice(), 1});<br>
torch::Tensor u = control(x1, x2, p1, p2, W);</p>
<pre><code>auto H = p1 * x2 + p2 * u * ((1 - x1 * x1) * x2 - x1) + W * u * u / 2 + 1;
return H;
</code></pre>
<p>}</p>
<p>TEST(HamiltonianTest, ppppppHTest)<br>
{<br>
torch::Tensor x = torch::tensor({0.5, -0.5}).view({1, -1}); // example tensor<br>
torch::Tensor p = torch::tensor({1.0, -1.0}).view({1, -1}); // example tensor<br>
double W = 1.0; // example parameter</p>
<pre><code>auto result = ppppppH&lt;double&gt;(x, p, W, hamiltonian);

// Print the result for verification
std::cout &lt;&lt; result &lt;&lt; std::endl;
</code></pre>
<p>}</p>
<p>Autograd loses track of the gradients in the second loop.  I have to set the grad_H_p_i manually.<br>
This should not be happening.</p>

username: panos
cooked: <p>My apologies for the formatting<br>
Here are the code snippets</p>
<pre><code class="lang-auto">template&lt;typename T&gt;
torch::Tensor ppppppH(const torch::Tensor &amp;x, 
                      const torch::Tensor &amp;p, 
                      T W,
                      std::function&lt;torch::Tensor(const torch::Tensor&amp;, const torch::Tensor&amp;, T)&gt; H)
{
    // Create tensors with gradient tracking for both x and p
    auto xt = x.clone();
    xt.detach_();
    xt=xt.set_requires_grad(false);
    auto pt = p.clone();
    pt.detach_();
    pt=pt.set_requires_grad(true);

    std::cerr &lt;&lt; "After setting gradient in ppppppH" &lt;&lt; std::endl;
    std::cerr &lt;&lt; "xt.requires_grad() = " &lt;&lt; xt.requires_grad() &lt;&lt; std::endl;
    std::cerr &lt;&lt; "pt.requires_grad() = " &lt;&lt; pt.requires_grad() &lt;&lt; std::endl;

    // Compute the Hamiltonian value
    auto Hvalue = H(xt, pt, W);
    std::cerr &lt;&lt; "Hvalue.requires_grad() = " &lt;&lt; Hvalue.requires_grad() &lt;&lt; std::endl;

    // Compute the first-order gradient of H with respect to p
    auto grad_H_p = torch::autograd::grad({Hvalue}, 
                                          {pt}, 
                                          {torch::ones_like(Hvalue)}, 
                                          true, 
                                          true, 
                                          true)[0];
    // Initialize a tensor to hold the third-order gradients
    std::cerr &lt;&lt; "grad_H_p.requires_grad() = " &lt;&lt; grad_H_p.requires_grad() &lt;&lt; std::endl;
    std::cerr &lt;&lt; "grad_H_p.grad_fn() = " &lt;&lt; grad_H_p.grad_fn() &lt;&lt; std::endl;
    std::cerr &lt;&lt; "grad_H_p shape = " &lt;&lt; grad_H_p.sizes() &lt;&lt; std::endl;
    std::cerr &lt;&lt; "grad_H_p values = " &lt;&lt; grad_H_p &lt;&lt; std::endl;

    auto batch_size = p.size(0);
    auto p_dim = p.size(1);
    torch::Tensor third_order_derivative = torch::zeros({batch_size, p_dim, p_dim, p_dim}, p.options());

    // Compute the second-order and third-order gradients
    for (int i = 0; i &lt; p_dim; ++i)
    {
        // Zero out the gradients of pt before computing the gradient for the next dimension
        if (pt.grad().defined()) {
            pt.grad().zero_();
        }

        // Compute the gradient of the i-th component of grad_H_p with respect to p
        auto grad_H_pati = grad_H_p.index({Slice(), i});
        auto grad_H_p_i = torch::autograd::grad({grad_H_pati}, 
                                                {pt}, 
                                                {torch::ones_like(grad_H_pati)}, 
                                                true, 
                                                true, 
                                                true)[0];
        std::cerr &lt;&lt; "grad_H_p_i.requires_grad() = " &lt;&lt; grad_H_p_i.requires_grad() &lt;&lt; std::endl;
        std::cerr &lt;&lt; "pt.requires_grad() = " &lt;&lt; pt.requires_grad() &lt;&lt; std::endl;
        for (int j = 0; j &lt; p_dim; ++j)
        {
            // Zero out the gradients of pt before computing the gradient for the next dimension
            if (pt.grad().defined()) {
                pt.grad().zero_();
            }
            //Check to see if grad_H_p_i lost the graph
            if ( grad_H_p_i.requires_grad() == false) {
                std::cerr &lt;&lt; "grad_H_p_i lost gradient tracking at this point." &lt;&lt; std::endl;
                grad_H_p_i.set_requires_grad(true);
            }


            // Compute the gradient of the j-th component of grad_H_p_i with respect to p
            auto grad_H_p_ij = torch::autograd::grad({grad_H_p_i.index({Slice(), j})}, 
                                                     {pt}, 
                                                     {torch::ones_like(grad_H_p_i.index({Slice(), j}))}, 
                                                     true, 
                                                     false, 
                                                     true)[0];
            std::cerr &lt;&lt; "grad_H_p_ij.requires_grad() = " &lt;&lt; grad_H_p_ij.requires_grad() &lt;&lt; std::endl;

            if (!grad_H_p_ij.requires_grad()) {
                std::cerr &lt;&lt; "grad_H_p_ij lost gradient tracking at this point." &lt;&lt; std::endl;
                continue;
            }

            //third_order_derivative.select(1, i).select(2, j).copy_(grad_H_p_ij);
            third_order_derivative.index_put_({Slice(), i, j}, grad_H_p_ij);
        }
    }

    // Return the third-order derivative tensor
    return third_order_derivative;
}

</code></pre>
<p>The code is tested using the following functions</p>
<pre><code class="lang-auto">torch::Tensor control(const torch::Tensor&amp; x1, 
                      const torch::Tensor&amp; x2,
                      const torch::Tensor&amp; p1,
                      const torch::Tensor&amp; p2, 
                      double W=1.0) {
    auto u = -p2 * ((1 - x1 * x1) * x2 - x1) / W;
    return u; // Return through copy elision
}

torch::Tensor hamiltonian(const torch::Tensor&amp; x, 
                          const torch::Tensor&amp; p, 
                          double W) {
    torch::Tensor p1 = p.index({Slice(), 0});  
    torch::Tensor p2 = p.index({Slice(), 1});  
    torch::Tensor x1 = x.index({Slice(), 0});  
    torch::Tensor x2 = x.index({Slice(), 1});  
    torch::Tensor u = control(x1, x2, p1, p2, W);

    auto H = p1 * x2 + p2 * u * ((1 - x1 * x1) * x2 - x1) + W * u * u / 2 + 1;
    return H;
}



TEST(HamiltonianTest, ppppppHTest)
{
    torch::Tensor x = torch::tensor({0.5, -0.5}).view({1, -1}); // example tensor
    torch::Tensor p = torch::tensor({1.0, -1.0}).view({1, -1}); // example tensor
    double W = 1.0; // example parameter

    auto result = ppppppH&lt;double&gt;(x, p, W, hamiltonian);

    // Print the result for verification
    std::cout &lt;&lt; result &lt;&lt; std::endl;

}

</code></pre>
<p>Autograd appears to lose track of the gradient in the inner loop.  This is easy to reproduce.  I believe the gradient should not have to be set manually like this.</p>