code: def test_auto_restore(self):
      tmp_dir = self.create_tempdir().full_path
      to_save = [
          CustomDC(foo=12, bar={"x": jnp.array((1, 4))}),
          np.array((2, 3)),
      ]
      target = [
          CustomDC(foo=0, bar={"x": jnp.array((0, 0))}),
          np.array((0, 0)),
      ]
      # Store an orbax ckpt
      config.update("flax_use_orbax_checkpointing", True)
      checkpoints.save_checkpoint(tmp_dir, to_save, 0, prefix="test_")
      # And a legacy ckpt
      config.update("flax_use_orbax_checkpointing", False)
      checkpoints.save_checkpoint(
          tmp_dir, to_save, 1, prefix="test_", keep=2
      )
  
      # Both gets restored with same API.
      restored = checkpoints.restore_checkpoint(
          os.path.join(tmp_dir, "test_0"), target=target
      )
      check_eq(restored, to_save)
      restored = checkpoints.restore_checkpoint(
          os.path.join(tmp_dir, "test_1"), target=target
      )
      check_eq(restored, to_save)
  
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 15 for calling the CheckpointsTest obj's underlying member method code with example usage: obj.test_auto_restore() without return types