tf.raw_ops.EnqueueTPUEmbeddingSparseTensorBatch

Eases the porting of code that uses tf.nn.embedding_lookup_sparse().

sample_indices[i], embedding_indices[i] and aggregation_weights[i] correspond to the ith feature. table_ids[i] indicates which embedding table to look up ith feature.

The tensors at corresponding positions in the three input lists (sample_indices, embedding_indices and aggregation_weights) must have the same shape, i.e. rank 1 with dim_size() equal to the total number of lookups into the table described by the corresponding feature.

Args
sample_indices A list of at least 1 Tensor objects with the same type in: int32, int64. A list of rank 1 Tensors specifying the training example to which the corresponding embedding_indices and aggregation_weights values belong. It corresponds to sp_ids.indices[:,0] in embedding_lookup_sparse().
embedding_indices A list with the same length as sample_indices of Tensor objects with the same type in: int32, int64. A list of rank 1 Tensors, indices into the embedding tables. It corresponds to sp_ids.values in embedding_lookup_sparse().
aggregation_weights A list with the same length as sample_indices of Tensor objects with the same type in: float32, float64. A list of rank 1 Tensors containing per training example aggregation weights. It corresponds to sp_weights.values in embedding_lookup_sparse().
mode_override A Tensor of type string. A string input that overrides the mode specified in the TPUEmbeddingConfiguration. Supported values are {'unspecified', 'inference', 'training', 'backward_pass_only'}. When set to 'unspecified', the mode set in TPUEmbeddingConfiguration is used, otherwise mode_override is used.
table_ids A list of ints. A list of integers specifying the identifier of the embedding table (offset of TableDescriptor in the TPUEmbeddingConfiguration) to lookup the corresponding input. The ith input is looked up using table_ids[i]. The size of the table_ids list must be equal to that of sample_indices, embedding_indices and aggregation_weights.
device_ordinal An optional int. Defaults to -1. The TPU device to use. Should be >= 0 and less than the number of TPU cores in the task on which the node is placed.
combiners An optional list of strings. Defaults to []. A list of string scalars, one for each embedding table that specify how to normalize the embedding activations after weighted summation. Supported combiners are 'mean', 'sum', or 'sqrtn'. It is invalid to have the sum of the weights be 0 for 'mean' or the sum of the squared weights be 0 for 'sqrtn'. If combiners isn't passed, the default is to use 'sum' for all tables.
max_sequence_lengths An optional list of ints. Defaults to [].
name A name for the operation (optional).
Returns
The created Operation.

© 2020 The TensorFlow Authors. All rights reserved.
Licensed under the Creative Commons Attribution License 3.0.
Code samples licensed under the Apache 2.0 License.
https://www.tensorflow.org/versions/r2.4/api_docs/python/tf/raw_ops/EnqueueTPUEmbeddingSparseTensorBatch