Check JAX version when using jax.layout.Format.

PiperOrigin-RevId: 778111564

Author: ChromeHearts

checkpoint/CHANGELOG.md

@@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [0.11.18] - 2025-07-01
+
+### Changed
+
+- For JAX>=0.6.2, JAX layout.Layout renamed to layout.Format
+
 ## [0.11.17] - 2025-06-30
 
 ### Added

checkpoint/orbax/checkpoint/_src/handlers/standard_checkpoint_handler_test_utils.py

@@ -42,7 +42,10 @@
   DLL = layout.Layout
 else:
   DLL = layout.DeviceLocalLayout  # type: ignore
-Format = layout.Format
+if jax.__version_info__ >= (0, 6, 2):
+  Format = layout.Format
+else:
+  Format = layout.Layout
 PyTree = Any
 SaveArgs = type_handlers.SaveArgs
 StandardRestoreArgs = standard_checkpoint_handler.StandardRestoreArgs
@@ -175,7 +178,7 @@ def test_custom_layout(self):
           if jax.__version_info__ >= (0, 6, 3)
           else arr.format.device_local_layout  # type: ignore
       )
-      custom_layout = Format(
+      custom_layout = Format(  # pytype: disable=wrong-keyword-args
           DLL(
               major_to_minor=arr_layout.major_to_minor[::-1],  # pytype: disable=attribute-error
               _tiling=arr_layout._tiling,  # pytype: disable=attribute-error

checkpoint/orbax/checkpoint/_src/serialization/serialization.py

@@ -42,7 +42,10 @@
 _CHECKPOINT_SUCCESS = 'checkpoint_write_success'
 
 Index = types.Index
-Format = layout.Format
+if jax.__version_info__ >= (0, 6, 2):
+  Format = layout.Format
+else:
+  Format = layout.Layout
 Shape = types.Shape
 
 
@@ -457,7 +460,7 @@ async def _read_array_index_and_device_put(
       sharding = jax.sharding.SingleDeviceSharding(
           device, memory_kind=memory_kind
       )
-      result.append(jax.device_put(shard, Format(dll, sharding)))
+      result.append(jax.device_put(shard, Format(dll, sharding)))  # pytype: disable=wrong-arg-types
   return result
 
 
@@ -508,7 +511,7 @@ async def read_and_create_array(
 
 
 async def async_deserialize(
-    user_sharding: jax.sharding.Sharding | Format,
+    user_sharding: jax.sharding.Sharding | Format,  # pytype: disable=wrong-arg-types  # pytype: disable=unsupported-operands
     tensorstore_spec: Union[ts.Spec, Dict[str, Any]],
     global_shape: Optional[Shape] = None,
     dtype: Optional[jnp.dtype] = None,

checkpoint/orbax/checkpoint/_src/serialization/serialization_test.py

@@ -41,8 +41,10 @@
   DLL = layout.Layout
 else:
   DLL = layout.DeviceLocalLayout  # type: ignore
-Format = layout.Format
-
+if jax.__version_info__ >= (0, 6, 2):
+  Format = layout.Format
+else:
+  Format = layout.Layout
 jax.config.update('jax_enable_x64', True)
 
 

checkpoint/orbax/checkpoint/_src/serialization/type_handlers.py

@@ -62,7 +62,10 @@
 PLACEHOLDER = ...
 PLACEHOLDER_TYPESTR = 'placeholder'
 
-Format = layout.Format
+if jax.__version_info__ >= (0, 6, 2):
+  Format = layout.Format
+else:
+  Format = layout.Layout
 Shape = arrays_types.Shape
 Scalar = Union[int, float, np.number]
 NamedSharding = jax.sharding.NamedSharding
@@ -812,9 +815,11 @@ class ArrayRestoreArgs(RestoreArgs):
   restore_type: Optional[Any] = jax.Array
   mesh: Optional[jax.sharding.Mesh] = None
   mesh_axes: Optional[jax.sharding.PartitionSpec] = None
-  sharding: Optional[Union[jax.sharding.Sharding, ShardingMetadata, Format]] = (
+  # pyformat: disable
+  sharding: Optional[Union[jax.sharding.Sharding, ShardingMetadata, Format]] = (  # type: ignore[invalid-annotation]
       None
   )
+  # pyformat: enable
   global_shape: Optional[Tuple[int, ...]] = None
   shape: Optional[Tuple[int, ...]] = None
   strict: bool = True

checkpoint/orbax/checkpoint/checkpoint_utils.py

@@ -35,7 +35,10 @@
 PyTree = Any
 STANDARD_ARRAY_TYPES = (int, float, np.ndarray, jax.Array)
 _SNAPSHOTS = '_SNAPSHOTS'
-Format = layout.Format
+if jax.__version_info__ >= (0, 6, 2):
+  Format = layout.Format
+else:
+  Format = layout.Layout
 PLACEHOLDER = type_handlers.PLACEHOLDER
 
 
@@ -428,7 +431,7 @@ def construct_restore_args(
 
   def _array_restore_args(
       value: Any,
-      sharding: Optional[jax.sharding.Sharding | Format],
+      sharding: Optional[jax.sharding.Sharding | Format],  # pytype: disable=unsupported-operands
       dtype: Optional[np.dtype] = None,
   ) -> type_handlers.ArrayRestoreArgs:
     return type_handlers.ArrayRestoreArgs(

checkpoint/orbax/checkpoint/experimental/v1/_src/serialization/array_leaf_handler.py

@@ -39,6 +39,11 @@
 ArrayDeserializationParam = types.DeserializationParam["AbstractArray"]
 Shape = arrays_types_v0.Shape
 
+if jax.__version_info__ >= (0, 6, 2):
+  Format = jax_layout.Format
+else:
+  Format = jax_layout.Layout
+
 
 class AbstractArray(Protocol):
   """Abstract representation of an array.
@@ -63,7 +68,7 @@ class AbstractArray(Protocol):
 
   shape: Shape | None
   dtype: jax.numpy.dtype | None
-  sharding: jax.sharding.Sharding | jax_layout.Format | None
+  sharding: jax.sharding.Sharding | Format | None  # pytype: disable=unsupported-operands
 
 
 @dataclasses.dataclass

checkpoint/orbax/checkpoint/version.py

@@ -17,7 +17,7 @@
 # A new PyPI release will be pushed everytime `__version__` is increased.
 # Also modify version and date in CHANGELOG.
 # LINT.IfChange
-__version__ = '0.11.17'
+__version__ = '0.11.18'
 # LINT.ThenChange(//depot//orbax/checkpoint/CHANGELOG.md)