jax

f941b14e - Merge pull request #37320 from samanklesaria:geometric_out_sharding