tf.compat.v1.nn.safe_embedding_lookup_sparse

Lookup embedding results, accounting for invalid IDs and empty features.

tf.compat.v1.nn.safe_embedding_lookup_sparse(
    embedding_weights,
    sparse_ids,
    sparse_weights=None,
    combiner='mean',
    default_id=None,
    name=None,
    partition_strategy='div',
    max_norm=None,
    allow_fast_lookup=False
)

The partitioned embedding in embedding_weights must all be the same shape except for the first dimension. The first dimension is allowed to vary as the vocabulary size is not necessarily a multiple of P. embedding_weights may be a PartitionedVariable as returned by using tf.compat.v1.get_variable() with a partitioner.

Invalid IDs (< 0) are pruned from input IDs and weights, as well as any IDs with non-positive weight. For an entry with no features, the embedding vector for default_id is returned, or the 0-vector if default_id is not supplied.

The ids and weights may be multi-dimensional SparseTensors or RaggedTensors with rank of 2. For SpareTensors with left-aligned non-zero entries which can be described as RaggedTensors, use of RaggedTensors can yield higher performance. Embeddings are always aggregated along the last dimension.

embedding_weights A single tensor representing the complete embedding tensor, or a list tensors all of same shape except for the first dimension, representing sharded embedding tensors. Alternatively, a PartitionedVariable, created by partitioning along dimension 0. Each element must be appropriately sized for the given partition_strategy.
sparse_ids SparseTensor of shape [d_0, d_1, ..., d_n] containing the ids, where d_0 is typically batch size, or a RaggedTensor with rank 2.
sparse_weights SparseTensor or RaggedTensor of same type and shape as sparse_ids, containing float weights corresponding to sparse_ids, or None if all weights are assumed to be 1.0.
combiner A string specifying how to combine embedding results for each entry. Currently "mean", "sqrtn" and "sum" are supported, with "mean" the default.
default_id The id to use for an entry with no features.
name A name for this operation (optional).
partition_strategy A string specifying the partitioning strategy. Currently "div" and "mod" are supported. Default is "div".
max_norm If not None, all embeddings are l2-normalized to max_norm before combining.
allow_fast_lookup An optional boolean specifying whether to allow simplified embedding lookups when params is a single tensor and max_norm is None. Setting this flag to True during training can cause the use of dense gradients with increased memory footprint.

A dense tensor representing the combined embeddings for the sparse ids. For each row in the dense tensor represented by sp_ids, the op looks up the embeddings for all ids in that row, multiplies them by the corresponding weight, and combines these embeddings as specified.

In other words, if

shape(combined embedding_weights) = [p0, p1, ..., pm]

and

shape(sparse_ids) = shape(sparse_weights) = [d0, d1, ..., dn]

then

shape(output) = [d0, d1, ... dn-1, p1, ..., pm].

For instance, if params is a 10x20 matrix, and sp_ids / sp_weights are

  [0, 0]: id 1, weight 2.0
  [0, 1]: id 3, weight 0.5
  [1, 0]: id -1, weight 1.0
  [2, 3]: id 1, weight 3.0

default_id is 0.

with combiner="mean", then the output will be a 3x20 matrix where

  output[0, :] = (params[1, :] * 2.0 + params[3, :] * 0.5) / (2.0 + 0.5)
  output[1, :] = (params[0, :] * 1.0) / 1.0
  output[2, :] = (params[1, :] * 3.0) / 3.0

ValueError if embedding_weights is empty.