PyTorch most efficient Jacobian/Hessian calculation
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)
&gt;&gt;&gt; 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)
&gt;&gt;&gt; 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?

functorch can speed up computations even more. E.g., this code is from the functorch docs for batched Jacobian calculation (Hessian works too):
batch_size = 64
Din = 31
Dout = 33

weight = torch.randn(Dout, Din)
print(f&quot;weight shape = {weight.shape}&quot;)
bias = torch.randn(Dout)

def predict(weight, bias, x):
    return F.linear(x, weight, bias).tanh()

x = torch.randn(batch_size, Din)
compute_batch_jacobian = vmap(jacrev(predict, argnums=2), in_dims=(None, None, 0))
batch_jacobian0 = compute_batch_jacobian(weight, bias, x)

