I have the following function

template
torch::Tensor ppppppH(const torch::Tensor &x,
const torch::Tensor &p,
T W,
std::function<torch::Tensor(const torch::Tensor&, const torch::Tensor&, T)> H)
{
// Create tensors with gradient tracking for both x and p
auto xt = x.clone();
xt.detach_();
auto pt = p.clone();
pt.detach_();

``````std::cerr << "After setting gradient in ppppppH" << std::endl;

// Compute the Hamiltonian value
auto Hvalue = H(xt, pt, W);

// Compute the first-order gradient of H with respect to p
{pt},
{torch::ones_like(Hvalue)},
true,
true,
true)[0];
// Initialize a tensor to hold the third-order gradients

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 < p_dim; ++i)
{
// Zero out the gradients of pt before computing the gradient for the next dimension
}

// Compute the gradient of the i-th component of grad_H_p with respect to p
{pt},
true,
true,
true)[0];
for (int j = 0; j < p_dim; ++j)
{
// Zero out the gradients of pt before computing the gradient for the next dimension
}
//Check to see if grad_H_p_i lost the graph
}

// Compute the gradient of the j-th component of grad_H_p_i with respect to p
{pt},
true,
false,
true)[0];

continue;
}

}
}

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

}

The function is tested using the following methods
torch::Tensor control(const torch::Tensor& x1,
const torch::Tensor& x2,
const torch::Tensor& p1,
const torch::Tensor& 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& x,
const torch::Tensor& 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<double>(x, p, W, hamiltonian);

// Print the result for verification
std::cout << result << std::endl;
``````

}

Autograd loses track of the gradients in the second loop. I have to set the grad_H_p_i manually.
This should not be happening.

My apologies for the formatting
Here are the code snippets

``````template<typename T>
torch::Tensor ppppppH(const torch::Tensor &x,
const torch::Tensor &p,
T W,
std::function<torch::Tensor(const torch::Tensor&, const torch::Tensor&, T)> H)
{
// Create tensors with gradient tracking for both x and p
auto xt = x.clone();
xt.detach_();
auto pt = p.clone();
pt.detach_();

std::cerr << "After setting gradient in ppppppH" << std::endl;

// Compute the Hamiltonian value
auto Hvalue = H(xt, pt, W);

// Compute the first-order gradient of H with respect to p
{pt},
{torch::ones_like(Hvalue)},
true,
true,
true)[0];
// Initialize a tensor to hold the third-order gradients

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 < p_dim; ++i)
{
// Zero out the gradients of pt before computing the gradient for the next dimension
}

// Compute the gradient of the i-th component of grad_H_p with respect to p
{pt},
true,
true,
true)[0];
for (int j = 0; j < p_dim; ++j)
{
// Zero out the gradients of pt before computing the gradient for the next dimension
}
//Check to see if grad_H_p_i lost the graph
}

// Compute the gradient of the j-th component of grad_H_p_i with respect to p
{pt},
true,
false,
true)[0];

continue;
}

}
}

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

``````

The code is tested using the following functions

``````torch::Tensor control(const torch::Tensor& x1,
const torch::Tensor& x2,
const torch::Tensor& p1,
const torch::Tensor& 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& x,
const torch::Tensor& 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<double>(x, p, W, hamiltonian);

// Print the result for verification
std::cout << result << std::endl;

}

``````

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.