code: def test_checkpointing_jax_array(self):
      global_mesh = jtu.create_global_mesh((4, 2), ("x", "y"))
      inp_shape = (8, 2)
      pspec = P("x", "y")
      num = math.prod(inp_shape)
  
      # First Array
      global_input_data1 = np.arange(num, dtype=np.int32).reshape(
          inp_shape
      )
      a1 = array.make_array_from_callback(
          inp_shape,
          NamedSharding(global_mesh, pspec),
          lambda idx: global_input_data1[idx],
      )
      ckpt_dir = pathlib.Path(self.create_tempdir("ckpt").full_path)
      ckpt_path1 = pathlib.Path(
          self.create_tempdir(f"{ckpt_dir}/first").full_path
      )
  
      # Second Array
      global_input_data2 = np.arange(
          num, num + num, dtype=np.int32
      ).reshape(inp_shape)
      a2 = array.make_array_from_callback(
          inp_shape,
          NamedSharding(global_mesh, pspec),
          lambda idx: global_input_data2[idx],
      )
      ckpt_path2 = pathlib.Path(
          self.create_tempdir(f"{ckpt_dir}/second").full_path
      )
  
      # Third Array
      def cb3(_):
          return np.array([], dtype=np.float32)
  
      global_mesh1d = jtu.create_global_mesh((8,), ("x",))
      a3 = array.make_array_from_callback(
          (0,), NamedSharding(global_mesh1d, P(None)), cb3
      )
      ckpt_path3 = pathlib.Path(
          self.create_tempdir(f"{ckpt_dir}/third").full_path
      )
  
      ckpt_paths = [str(ckpt_path1), str(ckpt_path2), str(ckpt_path3)]
      tspecs = jax.tree_util.tree_map(
          serialization.get_tensorstore_spec, ckpt_paths
      )
  
      manager = serialization.GlobalAsyncCheckpointManager()
      manager.serialize(
          [a1, a2, a3],
          tspecs,
          on_commit_callback=partial(
              self._on_commit_callback, ckpt_dir, ckpt_dir
          ),
      )
      manager.wait_until_finished()
  
      m1, m2, m3 = serialization.run_deserialization(
          [
              NamedSharding(global_mesh, pspec),
              NamedSharding(global_mesh, P("x")),
              NamedSharding(global_mesh1d, P(None)),
          ],
          tspecs,
      )
  
      self.assertIsInstance(m1, array.ArrayImpl)
      self.assertArraysEqual(
          np.asarray(m1.addressable_shards[0].data),
          np.array([[0], [2]], dtype=np.int32),
      )
      self.assertArraysEqual(
          np.asarray(m1.addressable_shards[1].data),
          np.array([[1], [3]], dtype=np.int32),
      )
      self.assertEqual(m1.addressable_shards[0].data.shape, (2, 1))
      self.assertEqual(m1.dtype, np.int32)
  
      self.assertIsInstance(m2, array.ArrayImpl)
      self.assertArraysEqual(
          np.asarray(m2.addressable_shards[0].data),
          np.array([[16, 17], [18, 19]], dtype=np.int32),
      )
      self.assertArraysEqual(
          np.asarray(m2.addressable_shards[1].data),
          np.array([[16, 17], [18, 19]], dtype=np.int32),
      )
      self.assertEqual(m2.addressable_shards[0].data.shape, (2, 2))
      self.assertEqual(m2.dtype, np.int32)
  
      self.assertIsInstance(m3, array.ArrayImpl)
      for i, s in enumerate(m3.addressable_shards):
          self.assertEqual(s.index, (slice(None),))
          self.assertEqual(s.replica_id, i)
          self.assertArraysEqual(
              np.asarray(s.data), np.array([], dtype=np.float32)
          )
      self.assertEqual(m3.dtype, np.float32)
  
imports: import asyncio
  import math
  from functools import partial
  import os
  import pathlib
  import tracemalloc
  from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  from jax._src import test_util
  from jax import config
  from jax._src import array
  from jax.sharding import NamedSharding, GSPMDSharding
  from jax.sharding import PartitionSpec
  from jax.experimental.array_serialization import serialization
  import numpy
  import tensorstore
  import unittest
label_desc: Use image node_id 4 for calling the CheckpointTest obj's underlying member method code with example usage: obj.test_checkpointing_jax_array() and returns: np