jax.experimental.custom_partitioning module .. automodule:: jax.experimental.custom_partitioning API .. autofunction:: custom_partitioning