From 5fdf2b2607c9b52c122ad81e60535cfaac825ac7 Mon Sep 17 00:00:00 2001
From: 8bitmp3 <19637339+8bitmp3@users.noreply.github.com>
Date: Tue, 13 Feb 2024 00:56:21 +0000
Subject: [PATCH] Add JAX mesh_utils docs

---
 docs/jax.experimental.mesh_utils.rst | 13 +++++++++++++
 docs/jax.experimental.rst            |  1 +
 2 files changed, 14 insertions(+)
 create mode 100644 docs/jax.experimental.mesh_utils.rst

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
 -----------------