EnqueueTPUEmbeddingIntegerBatch

public final class EnqueueTPUEmbeddingIntegerBatch

An op that enqueues a list of input batch tensors to TPUEmbedding.

Nested Classes

class EnqueueTPUEmbeddingIntegerBatch.Options Optional attributes for EnqueueTPUEmbeddingIntegerBatch  

Public Methods

static EnqueueTPUEmbeddingIntegerBatch
create(Scope scope, Iterable<Operand<Integer>> batch, Operand<String> modeOverride, Options... options)
Factory method to create a class wrapping a new EnqueueTPUEmbeddingIntegerBatch operation.
static EnqueueTPUEmbeddingIntegerBatch.Options
deviceOrdinal(Long deviceOrdinal)

Inherited Methods

Public Methods

public static EnqueueTPUEmbeddingIntegerBatch create (Scope scope, Iterable<Operand<Integer>> batch, Operand<String> modeOverride, Options... options)

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

Parameters
scope current scope
batch A list of 1D tensors, one for each embedding table, containing the indices into the tables.
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.
options carries optional attributes values
Returns
  • a new instance of EnqueueTPUEmbeddingIntegerBatch

public static EnqueueTPUEmbeddingIntegerBatch.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.