``jax.experimental.custom_partitioning`` module
===============================================

.. automodule:: jax.experimental.custom_partitioning

API
---

.. autofunction:: custom_partitioning