tfrs.experimental.layers.embedding.PartialTPUEmbedding

Partial TPU Embedding layer.

tfrs.experimental.layers.embedding.PartialTPUEmbedding(
    feature_config,
    optimizer: tf.keras.optimizers.Optimizer,
    pipeline_execution_with_tensor_core: bool = False,
    batch_size: Optional[int] = None,
    size_threshold: Optional[int] = 10000
) -> None

This layer is composed of tfrs.layers.embedding.TPUEmbedding and tf.keras.layers.Embedding embedding layers. When training on TPUs, it is preferable to use TPU Embedding layers for large tables (as they are sharded accross TPU cores) and Keras embedding layer for small tables. For tables with vocab sizes less than size_threshold a Keras embedding layer will be used, above that threshold a TPU embedding layer will be used.

This layer will be applied on a dictionary of feature_name, categorical_tensor pairs and return a dictionary of string-to-tensor of feature_name, embedded_value pairs.

feature_config A nested structure of tf.tpu.experimental.embedding.FeatureConfig configs.
optimizer An optimizer used for TPU embeddings.
pipeline_execution_with_tensor_core If True, the TPU embedding computations will overlap with the TensorCore computations (and hence will be one step old with potential correctness drawbacks). Set to True for improved performance.
batch_size If set, this will be used as the global batch size and override the autodetection of the batch size from the layer's input. This is necesarry if all inputs to the layer's call are SparseTensors.
size_threshold A threshold for table sizes below which a Keras embedding layer is used, and above which a TPU embedding layer is used. Set size_threshold=0 to use TPU embedding for all tables and size_threshold=None to use only Keras embeddings.

keras_embedding_layers Returns a dictionary mapping feature names to Keras embedding layers.
tpu_embedding Returns TPUEmbedding or None if only Keras embeddings are used.

Methods

call

View source

call(
    inputs: Dict[str, Tensor]
) -> Dict[str, tf.Tensor]

Computes the output of the embedding layer.

It expects a string-to-tensor (or SparseTensor/RaggedTensor) dict as input, and outputs a dictionary of string-to-tensor of feature_name, embedded_value pairs. Note that SparseTensor/RaggedTensor are only supported for TPUEmbedding and are not supported for Keras embeddings.

Args
inputs A string-to-tensor (or SparseTensor/RaggedTensor) dictionary.

Returns
output A dictionary of string-to-tensor of feature_name, embedded_value pairs.

Raises
ValueError if no tf.Tensor is passed to a Keras embedding layer.