jax.experimental.custom_partitioning module

.. automodule:: jax.experimental.custom_partitioning

API

.. autofunction:: custom_partitioning