flax

c1832a39 - Canonicalize PartitionSpec so that we can delete ParsedPartitionSpec

Canonicalize PartitionSpec so that we can delete ParsedPartitionSpec * `_partitions` is now canonicalized and only contains `tuples`, `None` or `UNCONSTRAINED`. * Cache the creating of sharding on ShapedArray since it's expensive to do it a lot of times * Change the `__hash__` and `__eq__` of `NamedSharding` to depend on `self.spec` instead of `self._parsed_pspec`. PiperOrigin-RevId: 730436599