code: def test_checkpoint_restore_before_variable_creation(self):
      self.skip_if_oss()
  
      class TestModule(module.Module):
          def __init__(self, initializer, rows):
              self._initializer = initializer
              self._rows = rows
  
              table = tpu_embedding_v2_utils.TableConfig(
                  vocabulary_size=self._rows,
                  dim=4,
                  initializer=self._initializer,
                  combiner="sum",
                  name="table",
              )
              feature_config = (
                  tpu_embedding_v2_utils.FeatureConfig(
                      table=table, name="feature"
                  ),
              )
              optimizer = tpu_embedding_v2_utils.SGD()
  
              self.tpu_embedding = tpu_embedding_v2.TPUEmbedding(
                  feature_config, optimizer
              )
  
          def create_embedding(self):
              # We aren't training so batch_size here doesn't matter.
              self.tpu_embedding.build(64)
  
      strategy = self._get_strategy()
      with strategy.scope():
          module1 = TestModule(
              init_ops_v2.Ones(), strategy.num_replicas_in_sync * 2
          )
          module1.create_embedding()
  
      checkpoint = util.Checkpoint(test_module=module1)
      checkpoint.save(self._get_tmpdir("restore_before_create", "save"))
  
      # Reinitialize the tpu
      strategy = self._get_strategy()
  
      with strategy.scope():
          module2 = TestModule(
              init_ops_v2.Zeros(), strategy.num_replicas_in_sync * 2
          )
  
      checkpoint = util.Checkpoint(test_module=module2)
      checkpoint.restore(
          self._get_tmpdir("restore_before_create", "save-1")
      )
  
      with strategy.scope():
          module2.create_embedding()
  
      def get_values(mid):
          return (
              mid._variables["table"]["parameters"].variables[0].numpy()
          )
  
      self.assertAllClose(
          np.ones((strategy.num_replicas_in_sync * 2, 4)),
          get_values(module2.tpu_embedding),
      )
  
      # Fetch the values from the TPU to check that they are the same.
      module2.tpu_embedding._retrieve_variables()
  
      self.assertAllClose(
          np.ones((strategy.num_replicas_in_sync * 2, 4)),
          get_values(module2.tpu_embedding),
      )
  
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 4 for calling the TPUEmbeddingCheckpointTest obj's underlying member method code with example usage: obj.test_checkpoint_restore_before_variable_creation() and returns: mid