Replace jax_helpers.batch_vmap with (eventual) upstreamed version of batch_vmap

[Jeremy Kun]