Eases the porting of code that uses tf.nn.embedding_lookup_sparse().
tf.raw_ops.EnqueueTPUEmbeddingArbitraryTensorBatch(
sample_indices_or_row_splits,
embedding_indices,
aggregation_weights,
mode_override,
device_ordinal=-1,
combiners=[],
name=None
)
embedding_indices[i] and aggregation_weights[i] correspond
to the 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_or_row_splits
|
A list of at least 1 Tensor objects with the same type in: int32 , int64 .
A list of rank 2 Tensors specifying the training example to which the
corresponding embedding_indices and aggregation_weights values belong.
If the size of its first dimension is 0, we assume each embedding_indices
belongs to a different sample. Both int32 and int64 are allowed and will
be converted to int32 internally.
Or a list of rank 1 Tensors specifying the row splits for splitting
embedding_indices and aggregation_weights into rows. It corresponds to
ids.row_splits in embedding_lookup(), when ids is a RaggedTensor. When
enqueuing N-D ragged tensor, only the last dimension is allowed to be ragged.
the row splits is 1-D dense tensor. When empty, we assume a dense tensor is
passed to the op Both int32 and int64 are allowed and will be converted to
int32 internally.
|
embedding_indices
|
A list with the same length as sample_indices_or_row_splits of Tensor objects with the same type in: int32 , int64 .
A list of rank 1 Tensors, indices into the embedding
tables. Both int32 and int64 are allowed and will be converted to
int32 internally.
|
aggregation_weights
|
A list with the same length as sample_indices_or_row_splits of Tensor objects with the same type in: float32 , float64 .
A list of rank 1 Tensors containing per training
example aggregation weights. Both float32 and float64 are allowed and will
be converted to float32 internally.
|
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.
|
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.
|
name
|
A name for the operation (optional).
|
Returns |
The created Operation.
|