class_docstr: Helper class for async checkpoint.
docstr: Handle TPUEmbedding.
  
  This is the only place where we populate object map in the class of
  `AsyncCheckpointHelper`. For all other checkpointable trackables, we
  populate object map using the trackable's own `_copy_trackable_to_cpu()`.
  
  Args:
    tpu_embedding: TPUEmbedding object to be handled.
  
  Raises:
    AttributeError: if the input trackable is not TPUEmbedding type.
docstr_tok: ["Handle","TPUEmbedding",".","This","is","the","only","place","where","we","populate","object","map","in","the","class","of","`","AsyncCheckpointHelper","`",".","For","all","other","checkpointable","trackables",",","we","populate","object","map","using","the","trackable","'s","own","`","_copy_trackable_to_cpu","(",")","`",".","Args",":","tpu_embedding",":","TPUEmbedding","object","to","be","handled",".","Raises",":","AttributeError",":","if","the","input","trackable","is","not","TPUEmbedding","type","."]
code: def _handle_tpu_embedding(self, tpu_embedding):
      """Handle TPUEmbedding.
  
      This is the only place where we populate object map in the class of
      `AsyncCheckpointHelper`. For all other checkpointable trackables, we
      populate object map using the trackable's own `_copy_trackable_to_cpu()`.
  
      Args:
        tpu_embedding: TPUEmbedding object to be handled.
  
      Raises:
        AttributeError: if the input trackable is not TPUEmbedding type.
      """
      if not hasattr(
          type(tpu_embedding), _TPU_EMBEDDING_ATTR
      ) or not callable(
          tpu_embedding._create_copy_for_async_checkpoint  # pylint: disable=protected-access
      ):
          raise AttributeError(
              "Expecting TPUEmbedding type; got %s"
              % type(tpu_embedding)
          )
  
      # Create a dummy TPUEmbedding object and add it to the object_map. This is
      # to prevent the TPUEmbedding's save_callback from being triggered because
      # the embedding values have already being retrieved by AsyncCheckpoint.
      # pylint: disable=protected-access
      new_embedding = tpu_embedding._create_copy_for_async_checkpoint(
          feature_config=tpu_embedding._feature_config,
          optimizer=tpu_embedding._table_config[0]
          if tpu_embedding._table_config
          else None,
          pipeline_execution_with_tensor_core=tpu_embedding._pipeline_execution_with_tensor_core,
      )
      self._object_map[tpu_embedding] = new_embedding
      # pylint: enable=protected-access
  
      if tpu_embedding not in self._tpu_embedding_objects:
          self._tpu_embedding_objects.append(tpu_embedding)
  
imports: import atexit
  import copy
  import queue
  import threading
  import time
  import weakref
  from absl import logging
  from tensorflow.python.checkpoint import checkpoint_context
  from tensorflow.python.checkpoint import trackable_view
  from tensorflow.python.distribute import device_util
  from tensorflow.python.eager import context
  from tensorflow.python.eager import def_function
  from tensorflow.python.eager import executor
  from tensorflow.python.framework import ops
  from tensorflow.python.ops import variables
  from tensorflow.python.saved_model.pywrap_saved_model import metrics
  from tensorflow.python.trackable import base
  from tensorflow.python.util import object_identity
label_desc: Use image node_id 8 for calling the AsyncCheckpointHelper obj's underlying member method code with example usage: obj._handle_tpu_embedding(tpu_embedding) without return types