jax

cd7ab0a9 - Changed to pmap_benchmark to make it runnable in Google (#2448)