jax.experimental.mesh_utils module

.. automodule:: jax.experimental.mesh_utils

API

.. autosummary:: :toctree: _autosummary

create_device_mesh create_hybrid_device_mesh