Add a custom __reduce__ for UnconstrainedSingleton because it can…

… be picked and then loaded back and we need the `id` of `P.UNCONSTRAINED` to match before and after loading.

PiperOrigin-RevId: 725874879

jax/_src/partition_spec.py

@@ -19,12 +19,18 @@ class UnconstrainedSingleton:
   def __repr__(self):
     return "UNCONSTRAINED"
 
+  def __reduce__(self):
+    return (_get_default_unconstrained, ())
+
 
 # Unconstrained sentinel value for PartitionSpec, representing a dimension for
 # which the user wants XLA to assign the best partitioning.
 # TODO(yashkatariya): May rename to AUTO.
 _UNCONSTRAINED_PARTITION = UnconstrainedSingleton()
 
+def _get_default_unconstrained():
+  return _UNCONSTRAINED_PARTITION
+
 
 class PartitionSpec(tuple):
   """Tuple describing how to partition an array across a mesh of devices.