implement vmap for jax solver#19

edgarWolf wants to merge 8 commits intomainedgarWolf/driftbench:mainfrom

18-speedup-jax-with-vmapedgarWolf/driftbench:18-speedup-jax-with-vmapCopy head branch name to clipboard