code: def test_keep(self, use_orbax, keep_every_n_steps):
      config.update("flax_use_orbax_checkpointing", use_orbax)
      tmp_dir = self.create_tempdir().full_path
      test_object = {"a": np.array([1, 2, 3], np.int32)}
      steps_start = 17
      steps_end = 37
      keep = 3
      increment = 5
  
      for step in range(steps_start, steps_end, increment):
          checkpoints.save_checkpoint(
              tmp_dir,
              test_object,
              step=step,
              keep=keep,
              keep_every_n_steps=keep_every_n_steps,
          )
  
      last_checkpoint = -float("inf")
      for step in range(steps_start, steps_end, increment):
          if ((steps_end - step) / increment <= keep) or (
              keep_every_n_steps
              and (step - last_checkpoint) >= keep_every_n_steps
          ):
              restored = checkpoints.restore_checkpoint(
                  tmp_dir, target=None, step=step
              )
              check_eq(restored, test_object)
              last_checkpoint = step
          else:
              with self.assertRaises(ValueError):
                  checkpoints.restore_checkpoint(
                      tmp_dir, target=None, step=step
                  )
  
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 6 for calling the CheckpointsTest obj's underlying member method code with example usage: obj.test_keep(use_orbax, keep_every_n_steps) without return types