rocm_jax/docs/jax.experimental.pjit.rst
Adam Paszke ee93ee221c Add a docstring for pjit
Co-authored-by: Skye Wanderman-Milne <skyewm@google.com>
PiperOrigin-RevId: 372393461
2021-05-06 12:00:49 -07:00

10 lines
130 B
ReStructuredText

jax.experimental.pjit module
============================
.. automodule:: jax.experimental.pjit
API
---
.. autofunction:: pjit