code: def save_to_path(ckpt_dir: str, state: CheckpointState):
      """Save the state in ckpt_dir."""
  
      if not os.path.exists(ckpt_dir):
          os.makedirs(ckpt_dir)
  
      is_numpy = lambda x: isinstance(x, (np.ndarray, jax.Array))
      flat_state = tree.flatten(state)
      nest_exemplar = tree.map_structure(is_numpy, state)
  
      array_path = os.path.join(ckpt_dir, _ARRAY_NAME)
      logging.info("Saving flattened array nest to %s", array_path)
  
      def _disabled_seek(*_):
          raise AttributeError("seek() is disabled on this object.")
  
      with open(array_path, "wb") as f:
          setattr(f, "seek", _disabled_seek)
          np.savez(f, *flat_state)
  
      exemplar_path = os.path.join(ckpt_dir, _EXEMPLAR_NAME)
      logging.info("Saving nest exemplar to %s", exemplar_path)
      with open(exemplar_path, "wb") as f:
          pickle.dump(nest_exemplar, f)
  
file_path: acme/acme/jax/savers.py
imports: import datetime
  import os
  import pickle
  from typing import Any
  from absl import logging
  from acme import core
  from acme.tf import savers
  import jax
  import numpy
  import tree
label_desc: Use image node_id 2 for calling a global function with example usage: save_to_path(ckpt_dir, state) without return types