jax.experimental.maps module .. automodule:: jax.experimental.maps API .. autofunction:: mesh .. autofunction:: xmap