jax

a0aa266a - Pass through the `use_shardy_partitioner` with `jax.config.jax_use_shardy_partitioner`.