Snapshot feature to save/load pathways checkpoint in host memory

PiperOrigin-RevId: 895163261

Author: Nikhil Bansal

checkpoint/orbax/checkpoint/experimental/v1/_src/testing/array_utils.py

@@ -43,19 +43,29 @@ def create_numpy_pytree(*, add: int = 0, include_scalars: bool = True):
 
 
 def create_sharded_pytree(
-    *, add: int = 0, reverse_devices: bool = False, include_scalars: bool = True
+    *,
+    add: int = 0,
+    reverse_devices: bool = False,
+    include_scalars: bool = True,
+    replicated_arrays: bool = False,
+    devices: list[jax.Device] | None = None,
+    slices: int = 2,
 ) -> tuple[tree_types.PyTree, tree_types.PyTree]:
   """Creates a sharded PyTree from `create_numpy_pytree`.
 
   Args:
     add: The value to add to leaf arrays.
     reverse_devices: Whether to reverse the devices in the mesh.
     include_scalars: Whether to include scalars in the pytree.
+    replicated_arrays: Whether to replicate arrays across devices.
+    devices: The devices to use for the mesh.
+    slices: The number of slices to use for the mesh.
 
   Returns:
     A tuple of (pytree, abstract_pytree).
   """
-  devices = jax.devices()
+  if devices is None:
+    devices = jax.devices()
   num_devices = len(devices)
   devices = (
       np.asarray(list(reversed(devices)))
@@ -64,13 +74,17 @@ def create_sharded_pytree(
   )
 
   mesh_2d = jax.sharding.Mesh(
-      devices.reshape((2, num_devices // 2)), ('x', 'y')
+      devices.reshape((slices, num_devices // slices)), ('x', 'y')
   )
   mesh_axes_2d = jax.sharding.PartitionSpec('x', 'y')
+  if replicated_arrays:
+    mesh_axes_2d = jax.sharding.PartitionSpec(None, 'y')
   mesh_1d = jax.sharding.Mesh(devices, ('x',))
   mesh_axes_1d = jax.sharding.PartitionSpec(
       'x',
   )
+  if replicated_arrays:
+    mesh_axes_1d = jax.sharding.PartitionSpec(None,)
   mesh_0d = jax.sharding.Mesh(devices, ('x',))
   mesh_axes_0d = jax.sharding.PartitionSpec(
       None,