code: def test_checkpoint_restore_loads(self):
      strategy = self._get_strategy()
      num_rows = strategy.num_replicas_in_sync
  
      def get_values(mid):
          return ops.convert_to_tensor(
              mid._variables["table"]["parameters"].variables[0]
          )
  
      with strategy.scope():
          first_mid_level_contents = np.ones((num_rows, 4))
          first_mid_level_optimizer = tpu_embedding_v2_utils.SGD(
              learning_rate=0.1
          )
          initializer = init_ops_v2.Constant(first_mid_level_contents)
  
          table = tpu_embedding_v2_utils.TableConfig(
              vocabulary_size=num_rows,
              dim=4,
              initializer=initializer,
              combiner="sum",
              name="table",
          )
          feature_config = (
              tpu_embedding_v2_utils.FeatureConfig(
                  table=table, name="feature"
              ),
          )
  
          first_mid_level = tpu_embedding_v2.TPUEmbedding(
              feature_config, first_mid_level_optimizer
          )
          first_mid_level.build(64)
  
      first_mid_level._load_variables()
  
      first_checkpoint = util.Checkpoint(model=first_mid_level)
      first_checkpoint.save(self._get_tmpdir("restore", "save"))
  
      tpu_cluster_resolver.initialize_tpu_system(self.resolver)
  
      with strategy.scope():
          second_mid_level_contents = np.ones((num_rows, 4)) * 2
          second_mid_level_optimizer = tpu_embedding_v2_utils.SGD(
              learning_rate=0.1
          )
          initializer = init_ops_v2.Constant(second_mid_level_contents)
  
          table = tpu_embedding_v2_utils.TableConfig(
              vocabulary_size=num_rows,
              dim=4,
              initializer=initializer,
              combiner="sum",
              name="table",
          )
          feature_config = (
              tpu_embedding_v2_utils.FeatureConfig(
                  table=table, name="feature"
              ),
          )
          second_mid_level = tpu_embedding_v2.TPUEmbedding(
              feature_config, second_mid_level_optimizer
          )
          second_mid_level.build(64)
  
      second_mid_level._load_variables()
  
      self.assertAllClose(
          second_mid_level_contents,
          get_values(second_mid_level),
          msg="Second mid level api should contain its initial values.",
      )
      # We restore the checkpoint of our first model into our second model.
      # This should load the first mid level API object onto the TPU.
      second_checkpoint = util.Checkpoint(model=second_mid_level)
      second_checkpoint.restore(self._get_tmpdir("restore", "save-1"))
  
      # Call retrieve here as a way to check what the TPU contains.
      # Calling the retrieve ops directly might make for a cleaner separation of
      # test and module, though.
      second_mid_level._retrieve_variables()
  
      self.assertAllClose(
          first_mid_level_contents,
          get_values(second_mid_level),
          msg="Second mid level api should have retrieved the first model values.",
      )
  
imports: from absl.testing import parameterized
  import numpy
  from tensorflow.python.checkpoint import checkpoint
  from tensorflow.python.compat import v2_compat
  from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver
  from tensorflow.python.eager import def_function
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import tensor_spec
  from tensorflow.python.module import module
  from tensorflow.python.ops import init_ops_v2
  from tensorflow.python.platform import test
  from tensorflow.python.saved_model import load
  from tensorflow.python.saved_model import save
  from tensorflow.python.tpu import tpu_embedding_for_serving
  from tensorflow.python.tpu import tpu_embedding_v2
  from tensorflow.python.tpu import tpu_embedding_v2_utils
  from tensorflow.python.tpu.tests import tpu_embedding_base_test
  from tensorflow.python.training import checkpoint_utils
label_desc: Use image node_id 3 for calling the TPUEmbeddingCheckpointTest obj's underlying member method code with example usage: obj.test_checkpoint_restore_loads() and returns: ops