code: def save(
      network,
      step,
      bucket,
      path,
      mp,
      aux=None,
      keep_n=3,
      delete_old=True,
  ):
      assert path
      client = storage.Client()
  
      if aux is None:
          aux = {}
  
      try:
          with open(f"gs://{bucket}/{path}/meta.json", "r") as f:
              meta = json.load(f)
      except:
          # create metadata file
          with open(f"gs://{bucket}/{path}/meta.json", "w") as f:
              json.dump({"step": 0, "checkpoints": [], "aux": {}}, f)
  
      # do sharded checkpoint writing
      start = time.time()
      res = []
      for shard_id in range(mp):
          write_ckpt(
              network.state,
              f"gs://{bucket}/{path}/step_{step}/",
              shard_id,
          )
  
      print(f"Wrote checkpoint in {time.time() - start:.06}s")
  
      with open(f"gs://{bucket}/{path}/meta.json", "r") as f:
          meta = json.load(f)
  
      meta["step"] = step
      meta["checkpoints"].append(step)
      all_aux = meta.get("aux", {})
  
      while len(meta["checkpoints"]) > keep_n:
          ckpt_to_delete = meta["checkpoints"].pop(0)
  
          try:
              del all_aux[str(ckpt_to_delete)]
          except:
              print(f"failed to delete the aux state for {step}")
  
          if delete_old:
              print(f"deleting checkpoint {ckpt_to_delete}")
              for blob in client.list_blobs(
                  bucket, prefix=f"{path}/step_{ckpt_to_delete}/"
              ):
                  # print(f"deleting {blob.name}")
                  assert path in blob.name
                  blob.delete()
          else:
              print(f"keeping checkpoint {ckpt_to_delete}")
  
      all_aux[step] = aux
      meta["aux"] = all_aux
  
      with open(f"gs://{bucket}/{path}/meta.json", "w") as f:
          json.dump(meta, f)
  
file_path: mesh-transformer-jax/device_train.py
imports: import argparse
  import json
  import time
  import jax
  import numpy
  import optax
  import wandb
  from tqdm import tqdm
  from mesh_transformer import util
  from mesh_transformer.checkpoint import read_ckpt, write_ckpt
  from mesh_transformer.transformer_shard import CausalTransformer
  from tfrecord_loader import TFRecordNewInputs
  from smart_open import open
  from google.cloud import storage
  from google.cloud.exceptions import NotFound
  from mesh_transformer.util import clip_by_global_norm, additive_weight_decay
label_desc: Use image node_id 2 for calling a global function with example usage: save(network, step, bucket, path, mp, aux, keep_n, delete_old) without return types