code: def test_save_restore_checkpoints(self, use_orbax):
      config.update("flax_use_orbax_checkpointing", use_orbax)
      tmp_dir = pathlib.Path(self.create_tempdir().full_path)
      test_object0 = {
          "a": np.array([0, 0, 0], np.int32),
          "b": np.array([0, 0, 0], np.int32),
      }
      test_object1 = {
          "a": np.array([1, 2, 3], np.int32),
          "b": np.array([1, 1, 1], np.int32),
      }
      test_object2 = {
          "a": np.array([4, 5, 6], np.int32),
          "b": np.array([2, 2, 2], np.int32),
      }
      new_object = checkpoints.restore_checkpoint(
          tmp_dir, test_object0, prefix="test_"
      )
      check_eq(new_object, test_object0)
      checkpoints.save_checkpoint(
          tmp_dir, test_object1, 0, prefix="test_", keep=1
      )
      self.assertIn("test_0", os.listdir(tmp_dir))
      new_object = checkpoints.restore_checkpoint(
          tmp_dir, test_object0, prefix="test_"
      )
      check_eq(new_object, test_object1)
      checkpoints.save_checkpoint(
          tmp_dir, test_object1, 1, prefix="test_", keep=1
      )
      checkpoints.save_checkpoint(
          tmp_dir, test_object2, 2, prefix="test_", keep=1
      )
      new_object = checkpoints.restore_checkpoint(
          tmp_dir, test_object0, prefix="test_"
      )
      check_eq(new_object, test_object2)
      checkpoints.save_checkpoint(
          tmp_dir, test_object2, 3, prefix="test_", keep=2
      )
      checkpoints.save_checkpoint(
          tmp_dir, test_object1, 4, prefix="test_", keep=2
      )
      new_object = checkpoints.restore_checkpoint(
          tmp_dir, test_object0, prefix="test_"
      )
      check_eq(new_object, test_object1)
      new_object = checkpoints.restore_checkpoint(
          tmp_dir, test_object0, step=3, prefix="test_"
      )
      check_eq(new_object, test_object2)
  
      # Restore with a specific checkpoint path, not the directory path.
      new_object = checkpoints.restore_checkpoint(
          os.path.join(tmp_dir, "test_3"), test_object0
      )
      check_eq(new_object, test_object2)
      # If a specific path is specified, but it does not exist, the same behavior
      # as when a directory is empty should apply: the target is returned
      # unchanged.
      new_object = checkpoints.restore_checkpoint(
          os.path.join(tmp_dir, "test_not_there"), test_object0
      )
      check_eq(new_object, test_object0)
      with self.assertRaises(ValueError):
          checkpoints.restore_checkpoint(
              tmp_dir, test_object0, step=5, prefix="test_"
          )
  
imports: import copy
  import os
  import pathlib
  from typing import Any
  import jax
  import numpy
  import orbax.checkpoint
  from absl.testing import absltest, parameterized
  from jax import numpy
  from flax import config, core, errors, io, struct
  from flax import linen
  from flax.training import checkpoints
label_desc: Use image node_id 4 for calling the CheckpointsTest obj's underlying member method code with example usage: obj.test_save_restore_checkpoints(use_orbax) without return types