jax

a77fac73 - Simpler integration tests for vmap and jvp