jax.experimental.shard_map module .. automodule:: jax.experimental.shard_map API .. autosummary:: :toctree: _autosummary shard_map