PyTorch most efficient Jacobian/Hessian calculation

Question

I am looking for the most efficient way to get the Jacobian of a function through Pytorch and have so far come up with the following solutions:

# Setup
def func(X):
    return torch.stack((X.pow(2).sum(1),
                        X.pow(3).sum(1),
                        X.pow(4).sum(1)),1)  

X = Variable(torch.ones(1,int(1e5))*2.00094, requires_grad=True).cuda()

# Solution 1:
t = time()
Y = func(X)
J = torch.zeros(3, int(1e5))

for i in range(3):
    J[i] = grad(Y[0][i], X, create_graph=True, retain_graph=True, allow_unused=True)[0]

print(time()-t)
>>> Output: 0.002 s

# Solution 2:
def Jacobian(f,X):
    X_batch = Variable(X.repeat(3,1), requires_grad=True)
    f(X_batch).backward(torch.eye(3).cuda(), retain_graph=True)
    return X_batch.grad

t = time()
J2 = Jacobian(func,X)
print(time()-t)
>>> Output: 0.001 s

Since there seem to be not a big difference between using a loop in the first solution than the second one, I wanted to ask if there might still be be a faster way to calculate a Jacobian in pytorch.

My other question is then also about what might be the most efficient way to calculate the Hessian.

Finally, does anyone know if something like this can be done easier or more efficient in TensorFlow?

Answer 1

The most efficient method is likely to use PyTorch's own inbuilt functions:

torch.autograd.functional.jacobian(func, x)
torch.autograd.functional.hessian(func, x)