code: def save_checkpoint(
      workdir: str,
      train_state: train_utils.TrainState,
      max_to_keep: int = 3,
      overwrite: bool = False,
  ):
      """Saves a checkpoint.
  
      First syncs the model state across replicas, then it unreplicates it by taking
      the train state of the first replica and saves it as a checkpoint.
  
      Args:
        workdir: Experiment directory for saving the checkpoint.
        train_state: An instance of TrainState that holds the state of training.
        max_to_keep: The number of checkpoints to keep.
        overwrite: Overwrite existing checkpoint  if a checkpoint
          at the current or a later step already exits (default: False).
      """
      if jax.process_index() == 0:
          # Get train state from the first replica.
          checkpoint_state = jax.device_get(
              jax_utils.unreplicate(train_state)
          )
          checkpoints.save_checkpoint(
              workdir,
              checkpoint_state,
              int(checkpoint_state.global_step),
              overwrite=overwrite,
              keep=max_to_keep,
          )
  
file_path: scenic/scenic/projects/ncr/utils.py
imports: import os
  from typing import Any, Dict, Optional, Tuple
  from flax import jax_utils
  from flax.training import checkpoints
  import jax
  import jax.numpy
  import numpy
  from scenic.train_lib import train_utils
  from tensorflow.io import gfile
label_desc: Use image node_id 1 for calling a global function with example usage: save_checkpoint(workdir, train_state, max_to_keep, overwrite) without return types