# tf.tpu.experimental.embedding.TPUEmbeddingForServing

<br />

|---|
| [![](https://www.tensorflow.org/images/GitHub-Mark-32px.png) View source on GitHub](https://github.com/tensorflow/tensorflow/blob/v2.13.1/tensorflow/python/tpu/tpu_embedding_for_serving.py#L39-L175) |

The TPUEmbedding mid level API running on CPU for serving.

#### View aliases


**Compat aliases for migration**

See
[Migration guide](https://www.tensorflow.org/guide/migrate) for
more details.

[`tf.compat.v1.tpu.experimental.embedding.TPUEmbeddingForServing`](https://www.tensorflow.org/api_docs/python/tf/tpu/experimental/embedding/TPUEmbeddingForServing)

<br />

    tf.tpu.experimental.embedding.TPUEmbeddingForServing(
        feature_config: Union[https://www.tensorflow.org/versions/r2.13/api_docs/python/tf/tpu/experimental/embedding/FeatureConfig, Iterable],
        optimizer: Optional[tpu_embedding_v2_utils._Optimizer]
    )

> [!NOTE]
> **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.

    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.


    # 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]
> **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.

<br />

<br />

<br />

| ## Args ||
|---|---|
| `feature_config` | A nested structure of [`tf.tpu.experimental.embedding.FeatureConfig`](https://www.tensorflow.org/versions/r2.13/api_docs/python/tf/tpu/experimental/embedding/FeatureConfig) configs. |
| `optimizer` | An instance of one of [`tf.tpu.experimental.embedding.SGD`](https://www.tensorflow.org/versions/r2.13/api_docs/python/tf/tpu/experimental/embedding/SGD), [`tf.tpu.experimental.embedding.Adagrad`](https://www.tensorflow.org/versions/r2.13/api_docs/python/tf/tpu/experimental/embedding/Adagrad) or [`tf.tpu.experimental.embedding.Adam`](https://www.tensorflow.org/versions/r2.13/api_docs/python/tf/tpu/experimental/embedding/Adam). When not created under TPUStrategy may be set to None to avoid the creation of the optimizer slot variables, useful for optimizing memory consumption when exporting the model for serving where slot variables aren't needed. |

<br />

<br />

<br />

<br />

| ## Raises ||
|---|---|
| `RuntimeError` | If created under TPUStrategy. |

<br />

<br />

<br />

<br />

| ## Attributes ||
|---|---|
| `embedding_tables` | Returns a dict of embedding tables, keyed by `TableConfig`. |

<br />

## Methods

### `build`

[View source](https://github.com/tensorflow/tensorflow/blob/v2.13.1/tensorflow/python/tpu/tpu_embedding_base.py#L128-L133)

    build()

Create variables and slots variables for TPU embeddings.

### `embedding_lookup`

[View source](https://github.com/tensorflow/tensorflow/blob/v2.13.1/tensorflow/python/tpu/tpu_embedding_for_serving.py#L154-L175)

    embedding_lookup(
        features: Any, weights: Optional[Any] = None
    ) -> Any

Apply standard lookup ops on CPU.

<br />

<br />

<br />

| Args ||
|---|---|
| `features` | A nested structure of [`tf.Tensor`](https://www.tensorflow.org/versions/r2.13/api_docs/python/tf/Tensor)s, [`tf.SparseTensor`](https://www.tensorflow.org/versions/r2.13/api_docs/python/tf/sparse/SparseTensor)s or [`tf.RaggedTensor`](https://www.tensorflow.org/versions/r2.13/api_docs/python/tf/RaggedTensor)s, with the same structure as `feature_config`. Inputs will be downcast to [`tf.int32`](https://www.tensorflow.org/versions/r2.13/api_docs/python/tf#int32). Only one type out of [`tf.SparseTensor`](https://www.tensorflow.org/versions/r2.13/api_docs/python/tf/sparse/SparseTensor) or [`tf.RaggedTensor`](https://www.tensorflow.org/versions/r2.13/api_docs/python/tf/RaggedTensor) is supported per call. |
| `weights` | If not `None`, a nested structure of [`tf.Tensor`](https://www.tensorflow.org/versions/r2.13/api_docs/python/tf/Tensor)s, [`tf.SparseTensor`](https://www.tensorflow.org/versions/r2.13/api_docs/python/tf/sparse/SparseTensor)s or [`tf.RaggedTensor`](https://www.tensorflow.org/versions/r2.13/api_docs/python/tf/RaggedTensor)s, matching the above, except that the tensors should be of float type (and they will be downcast to [`tf.float32`](https://www.tensorflow.org/versions/r2.13/api_docs/python/tf#float32)). For [`tf.SparseTensor`](https://www.tensorflow.org/versions/r2.13/api_docs/python/tf/sparse/SparseTensor)s we assume the `indices` are the same for the parallel entries from `features` and similarly for [`tf.RaggedTensor`](https://www.tensorflow.org/versions/r2.13/api_docs/python/tf/RaggedTensor)s we assume the row_splits are the same. |

<br />

<br />

<br />

<br />

| Returns ||
|---|---|
| A nested structure of Tensors with the same structure as input features. ||

<br />

### `__call__`

[View source](https://github.com/tensorflow/tensorflow/blob/v2.13.1/tensorflow/python/tpu/tpu_embedding_base.py#L135-L139)

    __call__(
        features: Any, weights: Optional[Any] = None
    ) -> Any

Call the mid level api to do embedding lookup.