mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-14 10:56:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add JAX mesh_utils docs

Browse Source
This commit is contained in:
8bitmp3 2024-02-13 00:56:21 +00:00
parent c1f234a95c
commit 5fdf2b2607
2 changed files with 14 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
docs/jax.experimental.mesh_utils.rst Normal file
View File

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

1
docs/jax.experimental.rst
View File

@ -25,6 +25,7 @@ Experimental Modules
jax.experimental.multihost_utils
jax.experimental.compilation_cache
jax.experimental.key_reuse
jax.experimental.mesh_utils
Experimental APIs
-----------------