Hi, I’m trying to create a Graph based model to learn on unstructured data using torch and torch_geometric in which my loss function will depend on 1st and 2nd order derivatives.
Within the model I use my points 3D coordinates to compute edge weights from the distances between them. The problem I’m having is that I need to compute a second order gradient w.r.t coordinates, I manage to obtain the first order gradient but not the second one. Here a minimal code to reproduce the issue:
import torch

# coordinates
x = torch.tensor( [[0.0,0.0,0.0],
[1.0,0.0,0.0],
[0.0,1.0,0.0],
[0.0,0.0,1.0],
[1.0,1.0,1.0]], requires_grad=True)

# Graphs edge connectivity
edges = torch.tensor(
[[0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4],
[1, 2, 3, 0, 2, 3, 4, 0, 1, 3, 4, 0, 1, 2, 4, 1, 2, 3]])

# Compute weight
row, col = edges
dist_sq = torch.sum((x[col] - x[row]).pow(2), dim=1)
weight = dist_sq.pow(-1)

# Compute 1st order derivative: This will have some non-null outputs
ones = torch.ones_like(weight)
qx = torch.autograd.grad(weight, x, grad_outputs=ones, create_graph=True)[0]
print(qx)

# Compute 2nd order derivative: Gradient will be null
# also a user-warning might be prompted regarding the use of .grad attribute
# on a non-leaf tensor which I tried to fix by forcing retain_grad() but nothing changed
ones = torch.ones_like(qx)
qxx = torch.autograd.grad(qx, x, grad_outputs=ones, create_graph=True)[0]
print(qxx)

If instead of computing the weight Tensor as an “edge attribute” I compute a “nodal” attribute by doing, say, weight = x.pow(2) I will get the proper gradients… but this is not what I’m looking for though… So far my only guess is that the shape change (from nodal to edge) might be posing a problem for the computation of the gradient.
Any ideas?
Thanks in advance!