tf.tpu.experimental.embedding.FeatureConfig

Configuration data for one embedding feature.

tf.tpu.experimental.embedding.FeatureConfig(
    table: tf.tpu.experimental.embedding.TableConfig,
    max_sequence_length: int = 0,
    validate_weights_and_indices: bool = True,
    output_shape: Optional[Union[List[int], tf.TensorShape]] = None,
    name: Optional[Text] = None
)

This class holds the configuration data for a single embedding feature. The main use is to assign features to tf.tpu.experimental.embedding.TableConfigs via the table parameter:

table_config_one = tf.tpu.experimental.embedding.TableConfig(
    vocabulary_size=...,
    dim=...)
table_config_two = tf.tpu.experimental.embedding.TableConfig(
    vocabulary_size=...,
    dim=...)
feature_config = {
    'feature_one': tf.tpu.experimental.embedding.FeatureConfig(
        table=table_config_one),
    'feature_two': tf.tpu.experimental.embedding.FeatureConfig(
        table=table_config_one),
    'feature_three': tf.tpu.experimental.embedding.FeatureConfig(
        table=table_config_two)}
embedding = tf.tpu.experimental.embedding.TPUEmbedding(
    feature_config=feature_config,
    batch_size=...
    optimizer=tf.tpu.experimental.embedding.Adam(0.1))

The above configuration has 2 tables, and three features. The first two features will be looked up in the first table and the third feature will be looked up in the second table.

You can also specify the output shape for each feature. The output shape should be the expected activation shape excluding the table dimension. For dense and sparse tensor, the output shape should be the same as the input shape excluding the last dimension. For ragged tensor, the output shape can mismatch the input shape.

When feeding features into embedding.enqueue they can be tf.Tensors, tf.SparseTensors or tf.RaggedTensors. When the argument max_sequence_length is 0, the default, you should expect a output of embedding.dequeue for this feature of shape (batch_size, dim). If max_sequence_length is greater than 0, the feature is embedded as a sequence and padded up to the given length. The shape of the output for this feature will be (batch_size, max_sequence_length, dim).

table An instance of tf.tpu.experimental.embedding.TableConfig, describing the table in which this feature should be looked up.
max_sequence_length If positive, the feature is a sequence feature with the corresponding maximum sequence length. If the sequence is longer than this, it will be truncated. If 0, the feature is not a sequence feature.
validate_weights_and_indices If true, uses safe_embedding_lookup during serving which ensures there are no empty rows and all weights and ids are positive at the expense of extra compute cost.
output_shape Optional argument to config the output shape of the feature activation. If provided, the feature feeding to the embedding.enqueue has to match the shape (for ragged tensor, the input shape and output shape can mismatch). If not provided, the shape can be either provided to the embedding.build or auto detected at the runtime.
name An optional string used to name the table. Must be defined if running on SparseCore.

ValueError if table is not an instance of tf.tpu.experimental.embedding.TableConfig.
ValueError if max_sequence_length not an integer or is negative.