jax

ab8720c2 - Check for existence of all `None`s in a spec to consider a sharding as replicated in reshape.

Check for existence of all `None`s in a spec to consider a sharding as replicated in reshape. This is because we want the logic for split/merge to activate even if a dimension is sharded on axis size `1` because that mesh axis can change the size if you run with a different mesh (see the parameterized test) so that we return the correct spec instead of just chopping off the last dimensions (when merging). Fixes https://github.com/jax-ml/jax/issues/29635 PiperOrigin-RevId: 774458015