|
|
@@ -13693,6 +13693,94 @@
|
|
|
"summary": "An op that enqueues a list of input batch tensors to TPUEmbedding."
|
|
|
}
|
|
|
},
|
|
|
+ {
|
|
|
+ "name": "EnqueueTPUEmbeddingRaggedTensorBatch",
|
|
|
+ "schema": {
|
|
|
+ "attributes": [
|
|
|
+ {
|
|
|
+ "default": {
|
|
|
+ "type": "type",
|
|
|
+ "value": 3
|
|
|
+ },
|
|
|
+ "description": "Must be one of the following: `int32`, `int64`.",
|
|
|
+ "name": "T1",
|
|
|
+ "type": "type"
|
|
|
+ },
|
|
|
+ {
|
|
|
+ "default": {
|
|
|
+ "type": "type",
|
|
|
+ "value": 3
|
|
|
+ },
|
|
|
+ "description": "Must be one of the following: `int32`, `int64`.",
|
|
|
+ "name": "T2",
|
|
|
+ "type": "type"
|
|
|
+ },
|
|
|
+ {
|
|
|
+ "default": {
|
|
|
+ "type": "type",
|
|
|
+ "value": 1
|
|
|
+ },
|
|
|
+ "description": "Must be one of the following: `float32`, `float64`.",
|
|
|
+ "name": "T3",
|
|
|
+ "type": "type"
|
|
|
+ },
|
|
|
+ {
|
|
|
+ "minimum": 1,
|
|
|
+ "name": "N",
|
|
|
+ "type": "int64"
|
|
|
+ },
|
|
|
+ {
|
|
|
+ "default": -1,
|
|
|
+ "description": "The TPU device to use. Should be >= 0 and less than the number\nof TPU cores in the task on which the node is placed.",
|
|
|
+ "name": "device_ordinal",
|
|
|
+ "type": "int64"
|
|
|
+ },
|
|
|
+ {
|
|
|
+ "default": [],
|
|
|
+ "description": "A list of string scalars, one for each embedding table that specify\nhow to normalize the embedding activations after weighted summation.\nSupported combiners are 'mean', 'sum', or 'sqrtn'. It is invalid to have\nthe sum of the weights be 0 for 'mean' or the sum of the squared weights be\n0 for 'sqrtn'. If combiners isn't passed, the default is to use 'sum' for\nall tables.",
|
|
|
+ "name": "combiners",
|
|
|
+ "type": "string[]"
|
|
|
+ },
|
|
|
+ {
|
|
|
+ "description": "A list of integers specifying the identifier of the embedding table\n(offset of TableDescriptor in the TPUEmbeddingConfiguration) to lookup the\ncorresponding input. The ith input is looked up using table_ids[i]. The size\nof the table_ids list must be equal to that of sample_indices,\nembedding_indices and aggregation_weights.",
|
|
|
+ "name": "table_ids",
|
|
|
+ "type": "int64[]"
|
|
|
+ },
|
|
|
+ {
|
|
|
+ "default": [],
|
|
|
+ "name": "max_sequence_lengths",
|
|
|
+ "type": "int64[]"
|
|
|
+ }
|
|
|
+ ],
|
|
|
+ "description": "sample_splits[i], embedding_indices[i] and aggregation_weights[i] correspond\nto the ith feature. table_ids[i] indicates which embedding table to look up ith\nfeature.\n\nThe tensors at corresponding positions in two of the input lists,\nembedding_indices and aggregation_weights, must have the same shape, i.e. rank 1\nwith dim_size() equal to the total number of lookups into the table described by\nthe corresponding feature.",
|
|
|
+ "inputs": [
|
|
|
+ {
|
|
|
+ "description": "A list of rank 1 Tensors specifying the break points for splitting\nembedding_indices and aggregation_weights into rows.\nIt corresponds to ids.row_splits in embedding_lookup(), when ids is a\nRaggedTensor.",
|
|
|
+ "name": "sample_splits",
|
|
|
+ "numberAttr": "N",
|
|
|
+ "typeAttr": "T1"
|
|
|
+ },
|
|
|
+ {
|
|
|
+ "description": "A list of rank 1 Tensors, indices into the embedding tables.\nIt corresponds to ids.values in embedding_lookup(), when ids is a RaggedTensor.",
|
|
|
+ "name": "embedding_indices",
|
|
|
+ "numberAttr": "N",
|
|
|
+ "typeAttr": "T2"
|
|
|
+ },
|
|
|
+ {
|
|
|
+ "description": "A list of rank 1 Tensors containing per training example\naggregation weights. It corresponds to the values field of a RaggedTensor\nwith the same row_splits as ids in embedding_lookup(), when ids is a\nRaggedTensor.",
|
|
|
+ "name": "aggregation_weights",
|
|
|
+ "numberAttr": "N",
|
|
|
+ "typeAttr": "T3"
|
|
|
+ },
|
|
|
+ {
|
|
|
+ "description": "A string input that overrides the mode specified in the\nTPUEmbeddingConfiguration. Supported values are {'unspecified', 'inference',\n'training', 'backward_pass_only'}. When set to 'unspecified', the mode set\nin TPUEmbeddingConfiguration is used, otherwise mode_override is used.",
|
|
|
+ "name": "mode_override",
|
|
|
+ "type": 7
|
|
|
+ }
|
|
|
+ ],
|
|
|
+ "summary": "Eases the porting of code that uses tf.nn.embedding_lookup()."
|
|
|
+ }
|
|
|
+ },
|
|
|
{
|
|
|
"name": "EnqueueTPUEmbeddingSparseBatch",
|
|
|
"schema": {
|
Lutz Roeder