Remove all usage of torch._vmap_internals._vmap and replace with new tools that came out in PyTorch 2. Specifically, all Jacobian and Hessian computation should now use the torch.func.jacfwd utility for forward-mode automatic differentiation