# tf.tpu.experimental.embedding.TPUEmbeddingV0

<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_v1.py#L37-L440) |

The TPUEmbedding mid level API running on TPU without Embedding accelerator.

#### View aliases


**Compat aliases for migration**

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

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

<br />

    tf.tpu.experimental.embedding.TPUEmbeddingV0(
        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 mid level API is not intended for large embedding table lookup. Embedding tables will be replicated across devices rather than sharding across them. To do large embedding table lookup, please use the [`tpu.experimental.embedding.TPUEmbedding`](https://www.tensorflow.org/versions/r2.13/api_docs/python/tf/tpu/experimental/embedding/TPUEmbedding) class. This class is an alternative way to do embedding lookups when the TPU doesn't support any version of embedding feature. See `tpu.experimental.tpu_hardware_feature.embedding_feature` for a detailed explanation.

This class has to be created under the `TPUStrategy`, Otherwise a RuntimeError
will be raised.

    strategy = tf.distribute.TPUStrategy(...)
    with strategy.scope():
      embedding = tf.tpu.experimental.embedding.TPUEmbeddingV0(
          feature_config=feature_config,
          optimizer=tf.tpu.experimental.embedding.SGD(0.1))

When creating a distributed dataset that is to be passed to the lookup
operation a special input option must be specified:

    distributed_dataset = (
        strategy.distribute_datasets_from_function(
            dataset_fn=...,
            options=tf.distribute.InputOptions(
                experimental_fetch_to_device=False))
    dataset_iterator = iter(distributed_dataset)

Below is an example of a training and evaluation step:

    optimizer = tf.keras.optimizers.SGD(0.1)

    @tf.function
    def training_step(dataset_iterator, num_steps):
      def tpu_step(embedding_features):
        with tf.GradientTape() as tape:
          tape.watch(embedding.embedding_table.values())
          activation = embedding(embedding_features)
          model_output = model(activations)
          loss = ...  # some function of labels and model_output

        embedding_gradients = tape.gradient(loss,
                                            embedding.embedding_table.values())
        optimizer.apply_gradients(list(zip(gradients,
                                  mid_level_api.embedding_tables.values())))
        # Insert your model gradient and optimizer application here

      for _ in tf.range(num_steps):
        strategy.run(tpu_step, args=(next(dataset_iterator), ))

    @tf.function
    def evalution_step(dataset_iterator, num_steps):
      def tpu_step(embedding_features):
        activations = embedding(embedding_features)
        model_output = model(activations)
        # Insert your evaluation code here.

      for _ in tf.range(num_steps):
        strategy.run(tpu_step, args=(next(dataset_iterator), ))

> [!NOTE]
> **Note:** The optimizer used here is a Keras optimizer. In order to make the slot variable creation stay consistent between Keras optimizers and embedding optimizers, the `slot_variable_creation_fn` argument of the embedding optimizers has to be passed with the Keras `add_slot` function. Also note that the slot names might be slightly different between them.

    optimizer = tf.keras.optimizers.Adagrad(learning_rate=0.1)

    def slot_variable_creation_fn(table, slot_names, slot_initializers):
        slots = {}
        for slot, initializer in zip(slot_names, slot_initializers):
          slots[slot] = optimizer.add_slot(table, slot, initializer)
        return slots

    embedding_optimizer = tf.experimental.embedding.Adagrad(
        learning_rate=0.1,
        slot_variable_creation_fn=slot_variable_creation_fn)

    # Use the embedding optimizer to create mid level api and keras optimizer to
    # apply gradients.

<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_v1.py#L239-L304)

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

Apply embedding lookup on TPUs using Tensorcore.

Note that all the sparse and ragged tensors will be converted to dense
tensors on CPU and then passed to the TPU to do embedding look up. Large
embedding lookup is not supported by this API, use the TPUEmbedding mid
level api instead.

<br />

<br />

<br />

| Args ||
|---|---|
| `features` | a nested structure of Tensors, SparseTensors or RaggedTensors. |
| `weights` | a nested structure of Tensors, SparseTensors or RaggedTensors or None for no weights. If not None, structure must match that of inputs, but entries are allowed to be None. |

<br />

<br />

<br />

<br />

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

<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.