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