code: def read_sharded_v2(state, dir, checkpoint_hosts, state_shard):
      files_per_host = checkpoint_hosts // jax.host_count()
  
      assert (
          files_per_host >= 1
      ), "can't restore model to larger pod than was trained on (yet)"
      assert (
          jax.host_count() * files_per_host == checkpoint_hosts
      ), "weird host count"
  
      if files_per_host == 1:
          head_print(
              "using fast path of checkpoint restore (save shards == read shards)"
          )
          parallel_read(state, dir + f"/shard_{jax.host_id()}.npz")
  
      @ray.remote
      def read_remote(old, fname):
          return parallel_read(old, fname, validate=False)
  
      start_idx = files_per_host * jax.host_id()
  
      skeleton = jax.tree_map(
          lambda x: jnp.zeros_like(x, shape=()), state
      )  # a full pytree just to carry dtypes
  
      refs = [
          read_remote.remote(skeleton, f"{dir}/shard_{i}.npz")
          for i in range(start_idx, start_idx + files_per_host)
      ]
  
      values = ray.get(refs)
  
      def all_array_equal(iterator):
          try:
              iterator = iter(iterator)
              first = next(iterator)
              return all(
                  jnp.array_equal(first, rest) for rest in iterator
              )
          except StopIteration:
              return True
  
      def reshard_v2(old, shard_strategy, *new_values):
          rep_dim_count = shard_strategy.count(None)
          total_dim_count = len(shard_strategy)
  
          # head_print("old.shape", old.shape)
          # head_print("shard_strategy", shard_strategy)
  
          assert len(old.shape) == total_dim_count
  
          if rep_dim_count == total_dim_count:
              # fully replicated
              assert all_array_equal(new_values)
              return fix_dtype(new_values[0])
  
          shard_dim = [
              idx
              for idx, dim in enumerate(shard_strategy)
              if dim is not None and "mp" in dim
          ]
  
          # only support sharding in 1d for now
          assert len(shard_dim) == 1
          shard_dim = shard_dim[0]
  
          ret_val = jnp.concatenate(
              fix_dtype(new_values), axis=shard_dim
          )
          assert old.shape == ret_val.shape
  
          return jax.device_put(ret_val, jax.devices("cpu")[0])
  
      # head_print("state", jax.tree_structure(state))
      # head_print("state_shard", jax.tree_structure(state_shard))
      # head_print("values", jax.tree_structure(values[0]))
  
      return jax.tree_multimap(
          reshard_v2, *([state, state_shard] + values)
      )
  
imports: import functools
  import io
  import json
  import time
  import jax
  import jax.numpy
  import numpy
  import multiprocessing
  import ray
  from smart_open import open
  from mesh_transformer.util import head_print
label_desc: Use image node_id 15 for calling a global function with example usage: read_sharded_v2(state, dir, checkpoint_hosts, state_shard) and returns: jax, parallel_read, jax, all, fix_dtype, True