tf.nn.embedding_lookup_sparse

TensorFlow 1 version View source on GitHub

Looks up embeddings for the given ids and weights from a list of tensors.

tf.nn.embedding_lookup_sparse(
    params, sp_ids, sp_weights, combiner=None, max_norm=None, name=None
)

This op assumes that there is at least one id for each row in the dense tensor represented by sp_ids (i.e. there are no rows with empty features), and that all the indices of sp_ids are in canonical row-major order.

sp_ids and sp_weights (if not None) are SparseTensors with rank of 2. Embeddings are always aggregated along the last dimension.

It also assumes that all id values lie in the range [0, p0), where p0 is the sum of the size of params along dimension 0.

If len(params) > 1, each element of sp_ids is partitioned between the elements of params according to the "div" partition strategy, which means we assign ids to partitions in a contiguous manner. For instance, 13 ids are split across 5 partitions as: [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10], [11, 12]].

If the id space does not evenly divide the number of partitions, each of the first (max_id + 1) % len(params) partitions will be assigned one more id.

params A single tensor representing the complete embedding tensor, or a list of tensors all of same shape except for the first dimension, representing sharded embedding tensors following "div" partition strategy.
sp_ids N x M SparseTensor of int64 ids where N is typically batch size and M is arbitrary.
sp_weights either a SparseTensor of float / double weights, or None to indicate all weights should be taken to be 1. If specified, sp_weights must have exactly the same shape and indices as sp_ids.
combiner A string specifying the reduction op. Currently "mean", "sqrtn" and "sum" are supported. "sum" computes the weighted sum of the embedding results for each row. "mean" is the weighted sum divided by the total weight. "sqrtn" is the weighted sum divided by the square root of the sum of the squares of the weights. Defaults to mean.
max_norm If not None, each embedding is clipped if its l2-norm is larger than this value, before combining.
name Optional name for the op.

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 params) = [p0, p1, ..., pm]

and

shape(sp_ids) = shape(sp_weights) = [d0, d1]

then

shape(output) = [d0, 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 0, weight 1.0
[2, 3]: id 1, weight 3.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

TypeError If sp_ids is not a SparseTensor, or if sp_weights is neither None nor SparseTensor.
ValueError If combiner is not one of {"mean", "sqrtn", "sum"}.