I have a neural network that's computing a vector quantity u. I'd like to compute first and second-order jacobians with respect to the input x, a single element.

Would anybody know how to do that in PyTorch? Below, the code snippet from my project:

import torch
import torch.nn as nnclass PINN(torch.nn.Module):def __init__(self, layers:list):super(PINN, self).__init__()self.linears = nn.ModuleList([])for i, dim in enumerate(layers[:-2]):self.linears.append(nn.Linear(dim, layers[i+1]))self.linears.append(nn.ReLU())self.linears.append(nn.Linear(layers[-2], layers[-1]))def forward(self, x):for layer in self.linears:x = layer(x)return x

I then instantiate my network:

n_in = 1
units = 50
q = 500pinn = PINN([n_in, units, units, units, q+1])
pinn

Which returns

PINN((linears): ModuleList((0): Linear(in_features=1, out_features=50, bias=True)(1): ReLU()(2): Linear(in_features=50, out_features=50, bias=True)(3): ReLU()(4): Linear(in_features=50, out_features=50, bias=True)(5): ReLU()(6): Linear(in_features=50, out_features=501, bias=True))
)

Then I compute both FO and SO jacobians

x = torch.randn(1, requires_grad=False)u_x = torch.autograd.functional.jacobian(pinn, x, create_graph=True)
print("First Order Jacobian du/dx of shape {}, and features\n{}".format(u_x.shape, u_x)u_xx = torch.autograd.functional.jacobian(lambda _: u_x, x)
print("Second Order Jacobian du_x/dx of shape {}, and features\n{}".format(u_xx.shape, u_xx)

Returns

First Order Jacobian du/dx of shape torch.Size([501, 1]), and features
tensor([[-0.0310],[ 0.0139],[-0.0081],[-0.0248],[-0.0033],[ 0.0013],[ 0.0040],[ 0.0273],...[-0.0197]], grad_fn=<ViewBackward>)

Second Order Jacobian du/dx of shape torch.Size([501, 1, 1]), and features
tensor([[[0.]],[[0.]],[[0.]],[[0.]],...[[0.]]])

Should not u_xx be a None vector if it didn't depend on x?

Thanks in advance

The second order Jacobian is known as the Hessian and can be computed easily using PyTorch's builtin functions:

torch.autograd.functional.hessian(func, inputs)