tf.tpu.experimental.embedding.CustomCombiner

Custom combiner for TPU embeddings.

tf.tpu.experimental.embedding.CustomCombiner(
    combiner_computation: tf.types.experimental.PolymorphicFunction,
    max_valency: int,
    num_weights: int = 0,
    initializer: Optional[init_ops_v2.Initializer] = None,
    combiner_weights_learning_rate: Union[float, Callable[[], float]] = 0.01
) -> Any

This class gives the user the ability to define a custom combiner for running embedding lookups on TPU with SparseCores.

The custom computation should be a function which takes the following arguments: (1) valency: an integer scalar that indicates the actual number of valent IDs in the current sample. (2) vectors: a 2D tensor of shape [max_valency, embedding_dim] that represents the embedding lookup results to be combined. The vectors are guaranteed to be in the same order as the embedding IDs appear in the input sample. (3) weights: this argument is only present if num_weights of this class is greater than 0. It is a 1D tensor of shape [num_weights] that will be back-propagated to during the backward pass. Currently only SGD optimizer is supported on these weights, with the learning rate being the combiner_weights_learning_rate argument of the constructor of this class.

The custom computation should return a 1D tensor of shape [embedding_dim] that represents the combined embedding vector. An example combiner computation is as follows (it simply sums over the valent embedding vectors and is semantically equivalent to the sum combiner):

@tf.function
def sum_combiner(valency, vectors):
  max_valency = vectors.shape[0]
  valid_mask = tf.range(max_valency) < valency
  vectors_masked = tf.where(
      tf.expand_dims(valid_mask, axis=-1),
      vectors,
      tf.zeros_like(vectors),
  )
  return tf.reduce_sum(vectors_masked, axis=0)

The custom computation is defined as a per-sample combiner function and it will be vectorized to the given batch size. This means certain constructs (such as control flow ops) may not be supported in the custom computation.

This class can be used in a tf.tpu.experimental.embedding.TableConfig as the combiner parameter to set a table specific combiner.

custom_combiner = tpu_embedding_v2_utils.CustomCombiner(
    sum_combiner,
    max_valency=16,
)

table_one = tf.tpu.experimental.embedding.TableConfig(
    vocabulary_size=...,
    dim=...,
    combiner=custom_combiner)
table_two = tf.tpu.experimental.embedding.TableConfig(
    vocabulary_size=...,
    dim=...)

feature_config = (
    tf.tpu.experimental.embedding.FeatureConfig(
        table=table_one),
    tf.tpu.experimental.embedding.FeatureConfig(
        table=table_two))

embedding = tf.tpu.experimental.embedding.TPUEmbedding(
    feature_config=feature_config,
    batch_size=...)

In this example, the combiner of the first table will be the sum_combiner. The second table will use the default mean combiner.

Methods

__eq__

View source

__eq__(
    other: Any
) -> Union[Any, bool]

Return self==value.

CUSTOM_COMBINER_SLOT_NAME 'custom_combiner_variables'