code: def parallel_read(old, fname, validate=True):
      old_vals, treedef = jax.tree_flatten(old)
  
      if "gs://" in fname:
          # TODO: make this actually parallel
          with open(fname, "rb") as f:
              buf = f.read()
              f_io = io.BytesIO(buf)
              loaded = np.load(f_io)
      else:
          loaded = np.load(fname, mmap_mode="r")
  
      new_vals = []
      for i in loaded:
          new_vals.append(loaded[i])
  
      assert len(new_vals) == len(old_vals), "Incompatible checkpoint"
  
      for o, n in zip(new_vals, old_vals):
          if validate:
              assert o.shape == n.shape, "Incompatible checkpoint"
  
      return jax.tree_unflatten(treedef, fix_dtype(new_vals))
  
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 11 for calling a global function with example usage: parallel_read(old, fname, validate) and returns: jax