jax

0bcf6e59 - Merge pull request #34809 from justinjfu:explicit_sharding