diff --git a/docs/jax.experimental.mesh_utils.rst b/docs/jax.experimental.mesh_utils.rst new file mode 100644 index 000000000..86de20d6b --- /dev/null +++ b/docs/jax.experimental.mesh_utils.rst @@ -0,0 +1,13 @@ +``jax.experimental.mesh_utils`` module +====================================== + +.. automodule:: jax.experimental.mesh_utils + +API +--- + +.. autosummary:: + :toctree: _autosummary + + create_device_mesh + create_hybrid_device_mesh diff --git a/docs/jax.experimental.rst b/docs/jax.experimental.rst index 15c907746..37cab679d 100644 --- a/docs/jax.experimental.rst +++ b/docs/jax.experimental.rst @@ -25,6 +25,7 @@ Experimental Modules jax.experimental.multihost_utils jax.experimental.compilation_cache jax.experimental.key_reuse + jax.experimental.mesh_utils Experimental APIs -----------------