I was actually able to solve it. This works exponentially faster than using the jacobian and hessian methods + diagonalization. However, I suspect this still isn’t at peak efficiency due to the for-loop. Hopefully something like this can be parallelized and implemented in the future.
# Creates first and second order jacobian for vector function f
def second_order_jacobians(f, wrt, create_graph=True):
    jacobian_1 = []
    jacobian_2 = []
    for i in range(f.shape[1]):
        j1 = grad(f[:,i:i+1], wrt, create_graph=create_graph, grad_outputs=torch.ones(f[:,i:i+1].shape).to(f.device))[0]
        j2 = grad(j1, wrt, create_graph=create_graph, grad_outputs=torch.ones(j1.shape).to(j1.device))[0]
        jacobian_1.append(j1)
        jacobian_2.append(j2)
    jacobian_1 = torch.stack(jacobian_1, -1)
    jacobian_2 = torch.stack(jacobian_2, -1)
    return jacobian_1, jacobian_2