Node-wise gradient of predictions wrt inputs for physics-informed graph neural networks

[aljubrmj]

Hi, I am working on a physics-informed graph neural network. For each node $n$, my physics loss term requires that I compute $\frac{\partial y_n}{\partial x_n}$, i.e., the gradient of the node-level prediction $y_n$ wrt the input node-level features $x_n$. Using torch.autograd.grad(y.sum(), x) rather results in a tensor whose elements represent $\sum\frac{\partial y_n}{\partial x_n}$, which is incorrect. I used a for loop to compute $\frac{\partial y_n}{\partial x_n}$ exactly, but it is too slow. See snippet below. I also tried with the PyTorch jacobian and vmap functions, but they are still slow given the size of graph I am dealing with. How can I achieve this computation efficiently? Is there a way to leverage the data graph structure?

dataset = Planetoid(root='.', name='Cora')
data = dataset.data
dataset.data.x = dataset.data.x[:,:num_features]
data.x = data.x.requires_grad_()
x, edge_index = data.x, data.edge_index
output = model(x, edge_index)

grads = []
for i in range(output.size(0)):
    grad = torch.autograd.grad(output[i], data.x, retain_graph=True, create_graph=True)[0][i]
    grads.append(grad)

grads = torch.stack(grads)

grads

Happy new year!