code: def write(x, ckpt_dir):
      # start = time.time()
      idx, i = x
      file_path = ckpt_dir + f"{idx}.npz"
      for _ in range(3):
          try:
              with open(file_path, "wb") as f:
                  np.savez(f, *i)
                  # cloudpickle.dump(i, f)
                  # print(f"written {idx} in {time.time() - start:.06}s")
              return
          except:
              print("save failed, trying again")
  
      print("save failed 3 times, exiting")
      raise Exception("save failed")
  
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 3 for calling a global function with example usage: write(x, ckpt_dir) without return types