Add support for jax.sharding.PositionalSharding in Orbax checkpoint…

…ing.

Joint work with yaning@

PiperOrigin-RevId: 658124242

checkpoint/orbax/checkpoint/metadata/sharding.py

@@ -21,6 +21,7 @@
 import logging
 from typing import List, Optional, Tuple, Union
 import jax
+from jax.experimental import mesh_utils
 import numpy as np
 
 PartitionSpecElement = Union[None, str, Tuple[str, ...]]
@@ -31,6 +32,8 @@
 _DEVICE_STR = 'device_str'
 _MESH_AXES = 'axis_names'
 _MESH_SHAPE = 'shape'
+_DEVICES_SHAPE = 'shape'
+_MEMORY_KIND = 'memory_kind'
 
 
 class ShardingTypes(enum.Enum):
@@ -220,7 +223,62 @@ class GSPMDShardingMetadata(ShardingMetadata):
 
 @dataclasses.dataclass
 class PositionalShardingMetadata(ShardingMetadata):
-  pass
+  """PositionalShardingMetadata representing `jax.sharding.PositionalSharding`."""
+  shape: np.ndarray
+  memory_kind: Optional[str] = None
+
+  @classmethod
+  def from_jax_sharding(
+      cls, jax_sharding: jax.sharding.PositionalSharding
+  ) -> 'PositionalShardingMetadata':
+    return cls(
+        shape=np.array(list(jax_sharding.shape)),
+        memory_kind=jax_sharding.memory_kind,
+    )
+
+  def to_jax_sharding(self) -> jax.sharding.PositionalSharding:
+    return jax.sharding.PositionalSharding(
+        mesh_utils.create_device_mesh(self.shape),
+        memory_kind=self.memory_kind,
+    )
+
+  @classmethod
+  def from_deserialized_dict(
+      cls, deserialized_dict: dict[str, str]
+  ) -> 'PositionalShardingMetadata':
+    if (
+        _DEVICES_SHAPE in deserialized_dict
+        and _MEMORY_KIND in deserialized_dict
+    ):
+      shape = np.array(deserialized_dict[_DEVICES_SHAPE])
+      memory_kind = deserialized_dict[_MEMORY_KIND]
+      return cls(
+          shape=shape,
+          memory_kind=memory_kind,
+      )
+    else:
+      raise ValueError(
+          f'Sharding data not found in deserialized_dict: {deserialized_dict}'
+      )
+
+  def to_serialized_string(self) -> str:
+    sharding_data = {}
+    sharding_data[_SHARDING_TYPE] = ShardingTypes.POSITIONAL_SHARDING.value
+    sharding_data[_DEVICES_SHAPE] = self.shape.tolist()
+    sharding_data[_MEMORY_KIND] = self.memory_kind
+    return json.dumps(sharding_data)
+
+  def __repr__(self):
+    return (
+        f'PositionalShardingMetadata(shape={self.shape},'
+        f' memory_kind={self.memory_kind})'
+    )
+
+  def __eq__(self, other):
+    return (
+        np.array_equal(self.shape, other.shape)
+        and self.memory_kind == other.memory_kind
+    )
 
 
 def from_jax_sharding(jax_sharding) -> Optional[ShardingMetadata]:
@@ -229,6 +287,8 @@ def from_jax_sharding(jax_sharding) -> Optional[ShardingMetadata]:
     return NamedShardingMetadata.from_jax_sharding(jax_sharding)
   elif isinstance(jax_sharding, jax.sharding.SingleDeviceSharding):
     return SingleDeviceShardingMetadata.from_jax_sharding(jax_sharding)
+  elif isinstance(jax_sharding, jax.sharding.PositionalSharding):
+    return PositionalShardingMetadata.from_jax_sharding(jax_sharding)
   else:
     logging.warning(
         'Conversion for %s has not been implemented.', type(jax_sharding)
@@ -247,6 +307,13 @@ def from_serialized_string(serialized_str) -> ShardingMetadata:
     return SingleDeviceShardingMetadata.from_deserialized_dict(
         deserialized_dict
     )
+  elif (
+      deserialized_dict[_SHARDING_TYPE]
+      == ShardingTypes.POSITIONAL_SHARDING.value
+  ):
+    return PositionalShardingMetadata.from_deserialized_dict(
+        deserialized_dict
+    )
   else:
     raise NotImplementedError(
         f'Conversion for {deserialized_dict[_SHARDING_TYPE]} has not been'

checkpoint/orbax/checkpoint/metadata/sharding_test.py

@@ -46,6 +46,39 @@ def test_convert_between_jax_named_sharding_and_sharding_metadata(self):
     self.assertIsInstance(converted_jax_sharding, jax.sharding.NamedSharding)
     self.assertEqual(converted_jax_sharding, jax_sharding)
 
+  def test_convert_between_jax_positional_sharding_and_sharding_metadata(
+      self,
+  ):
+    # Convert from `jax.sharding.PositionalSharding` to
+    # `PositionalShardingMetadata`
+    jax_sharding = jax.sharding.PositionalSharding(jax.devices()).reshape(
+        [1, -1]
+    )
+    expected_positional_sharding_metadata = (
+        sharding_metadata.PositionalShardingMetadata(jax_sharding.shape)
+    )
+    converted_positional_sharding_metadata = (
+        sharding_metadata.from_jax_sharding(jax_sharding)
+    )
+    self.assertIsInstance(
+        converted_positional_sharding_metadata,
+        sharding_metadata.PositionalShardingMetadata,
+    )
+    self.assertEqual(
+        converted_positional_sharding_metadata,
+        expected_positional_sharding_metadata,
+    )
+
+    # Convert from `PositionalShardingMetadata` to
+    # `jax.sharding.PositionalSharding`
+    converted_jax_sharding = (
+        converted_positional_sharding_metadata.to_jax_sharding()
+    )
+    self.assertIsInstance(
+        converted_jax_sharding, jax.sharding.PositionalSharding
+    )
+    self.assertEqual(converted_jax_sharding, jax_sharding)
+
   def test_convert_between_jax_single_device_sharding_and_sharding_metadata(
       self,
   ):
@@ -134,16 +167,6 @@ def test_single_device_sharding_string_to_metadata(
         single_device_sharding_metadata,
     )
 
-  def test_convert_to_jax_sharding_unsupported_types(self):
-    jax_sharding = jax.sharding.PositionalSharding(jax.devices())
-    warning_message = (
-        "Conversion for <class 'jax._src.sharding_impls.PositionalSharding'>"
-        " has not been implemented."
-    )
-    with self.assertLogs(level="WARNING") as log_output:
-      sharding_metadata.from_jax_sharding(jax_sharding)
-      self.assertEqual(log_output[0][0].message, warning_message)
-
 
 if __name__ == "__main__":
   absltest.main()