class_docstr: The TPUEmbedding mid level API running on CPU for serving.
  
  Note: This class is intended to be used for embedding tables that are trained
  on TPU and to be served on CPU. Therefore the class should be only initialized
  under non-TPU strategy. Otherwise an error will be raised.
  
  You can first train your model using the TPUEmbedding class and save the
  checkpoint. Then use this class to restore the checkpoint to do serving.
  
  First train a model and save the checkpoint.
  ```python
  model = model_fn(...)
  strategy = tf.distribute.TPUStrategy(...)
  with strategy.scope():
    embedding = tf.tpu.experimental.embedding.TPUEmbedding(
        feature_config=feature_config,
        optimizer=tf.tpu.experimental.embedding.SGD(0.1))
  
  # Your custom training code.
  
  checkpoint = tf.train.Checkpoint(model=model, embedding=embedding)
  checkpoint.save(...)
  
  ```
  
  Then restore the checkpoint and do serving.
  ```python
  
  # Restore the model on CPU.
  model = model_fn(...)
  embedding = tf.tpu.experimental.embedding.TPUEmbeddingForServing(
        feature_config=feature_config,
        optimizer=tf.tpu.experimental.embedding.SGD(0.1))
  
  checkpoint = tf.train.Checkpoint(model=model, embedding=embedding)
  checkpoint.restore(...)
  
  result = embedding(...)
  table = embedding.embedding_table
  ```
  
  NOTE: This class can also be used to do embedding training on CPU. But it
  requires the conversion between keras optimizer and embedding optimizers so
  that the slot variables can stay consistent between them.
class_docstr_tok: ["The","TPUEmbedding","mid","level","API","running","on","CPU","for","serving",".","Note",":","This","class","is","intended","to","be","used","for","embedding","tables","that","are","trained","on","TPU","and","to","be","served","on","CPU",".","Therefore","the","class","should","be","only","initialized","under","non-TPU","strategy",".","Otherwise","an","error","will","be","raised",".","You","can","first","train","your","model","using","the","TPUEmbedding","class","and","save","the","checkpoint",".","Then","use","this","class","to","restore","the","checkpoint","to","do","serving",".","First","train","a","model","and","save","the","checkpoint",".","``","`","python","model","=","model_fn","(","...",")","strategy","=","tf.distribute.TPUStrategy","(","...",")","with","strategy.scope","(",")",":","embedding","=","tf.tpu.experimental.embedding.TPUEmbedding","(","feature_config=feature_config",",","optimizer=tf.tpu.experimental.embedding.SGD","(","0.1",")",")","#","Your","custom","training","code",".","checkpoint","=","tf.train.Checkpoint","(","model=model",",","embedding=embedding",")","checkpoint.save","(","...",")","``","`","Then","restore","the","checkpoint","and","do","serving",".","``","`","python","#","Restore","the","model","on","CPU",".","model","=","model_fn","(","...",")","embedding","=","tf.tpu.experimental.embedding.TPUEmbeddingForServing","(","feature_config=feature_config",",","optimizer=tf.tpu.experimental.embedding.SGD","(","0.1",")",")","checkpoint","=","tf.train.Checkpoint","(","model=model",",","embedding=embedding",")","checkpoint.restore","(","...",")","result","=","embedding","(","...",")","table","=","embedding.embedding_table","``","`","NOTE",":","This","class","can","also","be","used","to","do","embedding","training","on","CPU",".","But","it","requires","the","conversion","between","keras","optimizer","and","embedding","optimizers","so","that","the","slot","variables","can","stay","consistent","between","them","."]
code: def _create_variables_from_stacked_tables(self):
      sc_layouts = sparse_core_layout_pb2.SparseCoreTableLayouts()
      sc_layouts.ParseFromString(
          self._get_sparse_core_table_layouts_str()
      )
      stacked_table_name_to_layouts = {}
      for layout in sc_layouts.tables:
          stacked_tables_list = (
              stacked_table_name_to_layouts.setdefault(
                  layout.stacked_table_name, []
              )
          )
          stacked_tables_list.append(layout)
      table_to_config = {
          table.name: table for table in self._table_config
      }
      variables = {}
      for (
          stacked_table_name,
          layouts,
      ) in stacked_table_name_to_layouts.items():
          logging.info(
              "Loading stacked table state variables(%s) for %s tables",
              stacked_table_name,
              len(layouts),
          )
          stacked_var_trackable = (
              tpu_embedding_v3_utils.SparseCoreStackedTableTrackable(
                  layouts, table_to_config
              )
          )
          # The stacked table is added as trackable to the embedding so that the
          # checkpoint key corresponsing to stacked table is read.
          self._track_trackable(
              stacked_var_trackable, stacked_table_name
          )
          variables.update(stacked_var_trackable.get_vars())
      return variables
  
imports: import functools
  from typing import Any, Dict, Iterable, Optional, Union
  from absl import logging
  from tensorflow.core.tpu.kernels import sparse_core_layout_pb2
  from tensorflow.python.distribute import distribute_lib
  from tensorflow.python.distribute import tpu_strategy
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import sparse_tensor
  from tensorflow.python.framework import tensor
  from tensorflow.python.framework.constant_op import constant
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import array_ops_stack
  from tensorflow.python.ops import embedding_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import sparse_ops
  from tensorflow.python.ops import variables
  from tensorflow.python.ops.ragged import ragged_tensor
  from tensorflow.python.tpu import tpu_embedding_base
  from tensorflow.python.tpu import tpu_embedding_v2_utils
  from tensorflow.python.tpu import tpu_embedding_v3_utils
  from tensorflow.python.types import core
  from tensorflow.python.util import nest
  from tensorflow.python.util.tf_export import tf_export
label_desc: Use image node_id 8 for calling the TPUEmbeddingForServing obj's underlying member method code with example usage: obj._create_variables_from_stacked_tables() and returns: variables