EnqueueTPUEmbeddingRaggedTensorBatch

public final class EnqueueTPUEmbeddingRaggedTensorBatch

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

sample_splits[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 two of the input lists, 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.

String OP_NAME The name of this op, as known by TensorFlow core engine

Public Methods

static EnqueueTPUEmbeddingRaggedTensorBatch.Options
combiners(List<String> combiners)
static <T extends TNumber, U extends TNumber, V extends TNumber> EnqueueTPUEmbeddingRaggedTensorBatch
create(Scope scope, Iterable<Operand<T>> sampleSplits, Iterable<Operand<U>> embeddingIndices, Iterable<Operand<V>> aggregationWeights, Operand<TString> modeOverride, List<Long> tableIds, Options... options)
Factory method to create a class wrapping a new EnqueueTPUEmbeddingRaggedTensorBatch operation.
static EnqueueTPUEmbeddingRaggedTensorBatch.Options
deviceOrdinal(Long deviceOrdinal)
static EnqueueTPUEmbeddingRaggedTensorBatch.Options
maxSequenceLengths(List<Long> maxSequenceLengths)

Inherited Methods

org.tensorflow.op.RawOp
final boolean
equals(Object obj)
final int
Operation
op()
Return this unit of computation as a single Operation.
final String
boolean
equals(Object arg0)
final Class<?>
getClass()
int
hashCode()
final void
notify()
final void
notifyAll()
String
toString()
final void
wait(long arg0, int arg1)
final void
wait(long arg0)
final void
wait()
org.tensorflow.op.Op
abstract ExecutionEnvironment
env()
Return the execution environment this op was created in.
abstract Operation
op()
Return this unit of computation as a single Operation.

Constants

public static final String OP_NAME

The name of this op, as known by TensorFlow core engine

Constant Value: "EnqueueTPUEmbeddingRaggedTensorBatch"

Public Methods

public static EnqueueTPUEmbeddingRaggedTensorBatch.Options combiners (List<String> combiners)

Parameters
combiners 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.

public static EnqueueTPUEmbeddingRaggedTensorBatch create (Scope scope, Iterable<Operand<T>> sampleSplits, Iterable<Operand<U>> embeddingIndices, Iterable<Operand<V>> aggregationWeights, Operand<TString> modeOverride, List<Long> tableIds, Options... options)

Factory method to create a class wrapping a new EnqueueTPUEmbeddingRaggedTensorBatch operation.

Parameters
scope current scope
sampleSplits A list of rank 1 Tensors specifying the break points for splitting embedding_indices and aggregation_weights into rows. It corresponds to ids.row_splits in embedding_lookup(), when ids is a RaggedTensor.
embeddingIndices A list of rank 1 Tensors, indices into the embedding tables. It corresponds to ids.values in embedding_lookup(), when ids is a RaggedTensor.
aggregationWeights A list of rank 1 Tensors containing per training example aggregation weights. It corresponds to the values field of a RaggedTensor with the same row_splits as ids in embedding_lookup(), when ids is a RaggedTensor.
modeOverride 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.
tableIds 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.
options carries optional attributes values
Returns
  • a new instance of EnqueueTPUEmbeddingRaggedTensorBatch

public static EnqueueTPUEmbeddingRaggedTensorBatch.Options deviceOrdinal (Long deviceOrdinal)

Parameters
deviceOrdinal 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.

public static EnqueueTPUEmbeddingRaggedTensorBatch.Options maxSequenceLengths (List<Long> maxSequenceLengths)