Add support for JAX batching via vmap

This will require updating the backend to iterate over the batch dimension, but that shouldn't be too terribly hard. Then, we'd need to add a simple batching function.

One question is how to interface batching with the `terms` interface.

https://jax.readthedocs.io/en/latest/notebooks/How_JAX_primitives_work.html#Batching