code: def test_checkpointing_with_bigger_shape_jax_array(self):
      global_mesh = jtu.create_global_mesh((2, 2), ("x", "y"))
      global_input_shape = (8, 2)
      num = math.prod(global_input_shape)
  
      global_input_data1 = np.arange(num, dtype=np.int32).reshape(
          global_input_shape
      )
  
      def cb1(index):
          return global_input_data1[index]
  
      arr = array.make_array_from_callback(
          global_input_shape,
          NamedSharding(global_mesh, P("x", "y")),
          cb1,
      )
      ckpt_dir = pathlib.Path(self.create_tempdir("first").full_path)
  
      ckpt_paths = [str(ckpt_dir)]
      tspecs = jax.tree_util.tree_map(
          serialization.get_tensorstore_spec, ckpt_paths
      )
  
      manager = serialization.GlobalAsyncCheckpointManager()
      manager.serialize(
          [arr],
          tspecs,
          on_commit_callback=partial(
              self._on_commit_callback, ckpt_dir, ckpt_dir
          ),
      )
      manager.wait_until_finished()
  
      ds = NamedSharding(
          jtu.create_global_mesh((4, 2), ("x", "y")), P("x", "y")
      )
  
      (m1,) = serialization.run_deserialization(
          [ds], tspecs, [(12, 2)], [np.float32]
      )
  
      expected_data = {
          0: np.array([[0], [2], [4]], dtype=np.float32),
          1: np.array([[1], [3], [5]], dtype=np.float32),
          2: np.array([[6], [8], [10]], dtype=np.float32),
          3: np.array([[7], [9], [11]], dtype=np.float32),
          4: np.array([[12], [14], [0]], dtype=np.float32),
          5: np.array([[13], [15], [0]], dtype=np.float32),
          6: np.array([[0], [0], [0]], dtype=np.float32),
          7: np.array([[0], [0], [0]], dtype=np.float32),
      }
  
      for l in m1.addressable_shards:
          self.assertArraysEqual(
              np.asarray(l.data), expected_data[l.device.id]
          )
  
      new_ds = GSPMDSharding.get_replicated(
          list(global_mesh.devices.flat)
      )
      (m2,) = serialization.run_deserialization(
          [new_ds], tspecs, [(8, 2)], [np.float32]
      )
      for l in m2.addressable_shards:
          self.assertArraysEqual(
              l.data, global_input_data1.astype("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 5 for calling the CheckpointTest obj's underlying member method code with example usage: obj.test_checkpointing_with_bigger_shape_jax_array() and returns: global_input_data1