Thanks, your solution works fine as long as one makes sure that input and output is 1D. I got an order of magnitude speed-up in the Jacobian calculation with your way.
import torch


def gradient(y, x, grad_outputs=None):
    """Compute dy/dx @ grad_outputs"""
    if grad_outputs is None:
        grad_outputs = torch.ones_like(y)
    grad = torch.autograd.grad(y, [x], grad_outputs = grad_outputs, create_graph=True)[0]
    return grad


def jacobian(y, x):
    """Compute dy/dx = dy/dx @ grad_outputs; 
    for grad_outputs in [1, 0, ..., 0], [0, 1, 0, ..., 0], ...., [0, ..., 0, 1]"""
    jac = torch.zeros(y.shape[0], x.shape[0])
    grad_outputs = torch.zeros_like(y)
    for i in range(y.shape[0]):
        grad_outputs[i] = 1
        jac[i] = gradient(y, x, grad_outputs = grad_outputs)
        grad_outputs[i] = 0
    return jac


x = torch.rand(3, requires_grad=True)
net = torch.nn.Linear(3, 2)
y = net(x)

jac1 = torch.autograd.functional.jacobian(net, x, create_graph=True)
jac2 = jacobian(y, x)

print(torch.allclose(jac1, jac2)) # True