code: def test_checkpoint_manager(self, shard_spec):
      tensor_a = stateless_random_ops.stateless_random_uniform(
          shape=[8, 16], seed=[0, 1]
      )
      tensor_b = stateless_random_ops.stateless_random_uniform(
          shape=[4, 4], seed=[0, 1]
      )
  
      layout = Layout(shard_spec, self.mesh)
  
      dvariable_a = d_variable.DVariable(
          numpy_util.pack_numpy(tensor_a, layout)
      )
      dvariable_b = d_variable.DVariable(
          numpy_util.pack_numpy(tensor_b, layout)
      )
  
      # Record a checkpoint with two dvariables.
      ckpt = checkpoint.Checkpoint(a=dvariable_a, b=dvariable_b)
  
      checkpoint_manager = checkpoint_management.CheckpointManager(
          ckpt, self.get_temp_dir(), max_to_keep=None
      )
  
      saved_path = checkpoint_manager.save()
  
      # Zero out the values of the DVariables so that we can restore
      # and check that the values are restored to the initial random values.
      dvariable_a.assign(
          numpy_util.pack_numpy(
              array_ops.zeros([8, 16], dtype=dtypes.float32), layout
          )
      )
      dvariable_b.assign(
          numpy_util.pack_numpy(
              array_ops.zeros([4, 4], dtype=dtypes.float32), layout
          )
      )
  
      ckpt.restore(saved_path)
  
      self.assertDTensorEqual(
          tensor_a, layout, dvariable_a.read_value()
      )
      self.assertDTensorEqual(
          tensor_b, layout, dvariable_b.read_value()
      )
  
imports: import gc
  from absl.testing import parameterized
  import numpy
  from tensorflow.dtensor.python import api
  from tensorflow.dtensor.python import d_variable
  from tensorflow.dtensor.python import layout
  from tensorflow.dtensor.python import numpy_util
  from tensorflow.dtensor.python.tests import test_util
  from tensorflow.python.checkpoint import checkpoint
  from tensorflow.python.checkpoint import checkpoint_management
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.module import module
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import stateless_random_ops
  from tensorflow.python.platform import test
label_desc: Use image node_id 4 for calling the DTensorSaveRestoreV2Test obj's underlying member method code with example usage: obj.test_checkpoint_manager(shard_spec) without return types