code: def load_ckpt_v2(model_state, dir, state_shard, load_opt):
      start = time.time()
      with open(dir + "meta.json", "r") as f:
          meta = json.load(f)
  
      ckpt_hosts = meta["total_hosts"]
  
      head_print(f"meta loaded in {time.time() - start:.06}s")
  
      new_state = {
          "step": np.array([meta["step"]]),
      }
  
      start = time.time()
      new_state["params"] = read_sharded_v2(
          model_state["params"],
          dir + "params",
          ckpt_hosts,
          state_shard["params"],
      )
      head_print(f"params loaded in {time.time() - start:.06}s")
  
      if not load_opt:
          return new_state
  
      start = time.time()
      new_state["opt_state"] = read_sharded_v2(
          model_state["opt_state"],
          dir + "opt_state",
          ckpt_hosts,
          state_shard["opt_state"],
      )
      head_print(f"opt_state loaded in {time.time() - start:.06}s")
  
      return new_state
  
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 16 for calling a global function with example usage: load_ckpt_v2(model_state, dir, state_shard, load_opt) and returns: new_state, new_state