Solved by Adam_Conkey in post #3 
              
              
                Thanks for the suggestion. It turns out we were setting vectorize=True in the hessian function which sounds like it’s doing strange things with vmap and it’s indicated it’s an experimental feature. Once I changed it back to False all was good.