jax

c0ba3626 - Merge pull request #26377 from mattjj:maintain-mutable-array-sharding