Computing a Hessian is expected to be slow I’m afraid.
Note that we have a function that computes the hessian directly in some cases that might help. But I don’t expect it to be much faster.

However, this method only works for scalar output functions

That’s the trick. If your function has a single input then the full Jacobian is recovered with a single backward pass and all the first order gradients become very easy to compute!