public final class EnqueueTPUEmbeddingIntegerBatch extends PrimitiveOp
Modifier and Type | Class and Description |
---|---|
static class |
EnqueueTPUEmbeddingIntegerBatch.Options
Optional attributes for
EnqueueTPUEmbeddingIntegerBatch |
operation
Modifier and Type | Method and Description |
---|---|
static EnqueueTPUEmbeddingIntegerBatch |
create(Scope scope,
Iterable<Operand<Integer>> batch,
Operand<String> modeOverride,
EnqueueTPUEmbeddingIntegerBatch.Options... options)
Factory method to create a class wrapping a new EnqueueTPUEmbeddingIntegerBatch operation.
|
static EnqueueTPUEmbeddingIntegerBatch.Options |
deviceOrdinal(Long deviceOrdinal) |
equals, hashCode, op, toString
public static EnqueueTPUEmbeddingIntegerBatch create(Scope scope, Iterable<Operand<Integer>> batch, Operand<String> modeOverride, EnqueueTPUEmbeddingIntegerBatch.Options... options)
scope
- current scopebatch
- 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 valuespublic static EnqueueTPUEmbeddingIntegerBatch.Options deviceOrdinal(Long deviceOrdinal)
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.Copyright © 2022. All rights reserved.