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

Author: yashk2810

flax/linen/spmd.py

@@ -105,7 +105,14 @@ def _logical_to_mesh_axes(
   # We assign mesh axes using a priority based ruleset over logical axis names.
   result: list[_UnassignedAxis | None | str | tuple[str, ...]]
   result = [
-      (_unassigned_axis if isinstance(name, str) else name)
+      (
+          _unassigned_axis
+          if (
+              isinstance(name, str)
+              or (isinstance(name, tuple) and len(name) == 1)
+          )
+          else name
+      )
       for name in array_dim_names
   ]
   for rule_model_name, rule_mesh_names in rules: