Skip to content

Commit

Permalink
Update tf-metadata.json
Browse files Browse the repository at this point in the history
  • Loading branch information
@lutzroeder
lutzroeder committed Mar 21, 2024
1 parent 7b46890 commit 26e138d
Showing 1 changed file with 236 additions and 0 deletions.
236 changes: 236 additions & 0 deletions source/tf-metadata.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9192,6 +9192,40 @@
}
]
},
{
"name": "ComputeDedupDataSizeV2",
"summary": "An op computes the size of the deduplication data from embedding core and returns the updated config.",
"description": "This op is to compute size of the deduplication data so to provide this\ninformation to the op that computes the tuple mask of deduplication data can\nhave static output shape.",
"attributes": [
{
"name": "config",
"type": "string",
"description": "Serialized TPUEmbeddingConfiguration proto."
},
{
"name": "embedding_partitions",
"type": "string",
"description": "Serialized EmbeddingPartitionsProto proto."
},
{
"name": "hbm_buffers_config",
"type": "string",
"description": "Serialized HbmBuffersConfig proto."
},
{
"name": "tpu_topology",
"type": "string",
"description": "Serialized TpuTopologyArgsProto proto."
}
],
"outputs": [
{
"name": "num_elements",
"description": "The size of the deduplicated data from infeed.",
"type": 3
}
]
},
{
"name": "ComputeDedupDataTupleMask",
"summary": "An op computes tuple mask of deduplication data from embedding core.",
Expand All @@ -9211,6 +9245,40 @@
}
]
},
{
"name": "ComputeDedupDataTupleMaskV2",
"summary": "An op computes tuple mask of deduplication data from embedding core.",
"description": "The deduplication data receiving from embedding core is a Tensor with\ntype=DT_VARIANT. The tensor itself is an XLA nested tuple, whose elements are\nrank 1 tensors. This op is to represents types and length of these elements.",
"attributes": [
{
"name": "config",
"type": "string",
"description": "Serialized TPUEmbeddingConfiguration proto."
},
{
"name": "embedding_partitions",
"type": "string",
"description": "Serialized EmbeddingPartitionsProto proto."
},
{
"name": "hbm_buffers_config",
"type": "string",
"description": "Serialized HbmBuffersConfig proto."
},
{
"name": "tpu_topology",
"type": "string",
"description": "Serialized TpuTopologyArgsProto proto."
}
],
"outputs": [
{
"name": "output_shape",
"description": "A 2-D int tensor represent mask of deduplication data tuple generated by\n`XlaRecvTPUEmbeddingDeduplicationData`. The tuple has several integer and float\ntype 1-D tensor tuple elements. The first dimenion of this output_shape 2-D\ntensor is tensor type of tuple elements, `0` represents integer tensor, `1`\nrepresents float tensor. The second dimension of `output_shape` gives length of\neach tuple element.",
"type": 3
}
]
},
{
"name": "Concat",
"category": "Tensor",
Expand Down Expand Up @@ -18928,6 +18996,34 @@
}
]
},
{
"name": "FinalizeTPUEmbeddingV2",
"summary": "An op that finalizes the TPUEmbedding configuration.",
"inputs": [
{
"name": "common_config",
"description": "A string-encoded common configuration proto containing metadata\nabout the TPUEmbedding partitioner output and the HBM size (in bytes) required\nfor operation.",
"type": 7
},
{
"name": "memory_config",
"description": "A string-encoded memory config proto containing metadata about\nthe memory allocations reserved for TPUEmbedding.",
"type": 7
}
],
"outputs": [
{
"name": "embedding_partitions",
"description": "A string-encoded embedding partitions proto describing how embedding tables are\npartitioned along their feature and ID.",
"type": 7
},
{
"name": "hbm_buffers_config",
"description": "A string-encoded HBM buffers config proto specifies where HBM buffers are\nlocated.",
"type": 7
}
]
},
{
"name": "Fingerprint",
"summary": "Generates fingerprint values.",
Expand Down Expand Up @@ -64214,6 +64310,54 @@
}
]
},
{
"name": "XlaRecvTPUEmbeddingActivationsV2",
"summary": "An op that receives embedding activations on the TPU.",
"description": "The TPU system performs the embedding lookups and aggregations. The results of\nthese aggregations are visible to the Tensorflow Graph as the outputs of a\nXlaRecvTPUEmbeddingActivations Op. This op returns a list containing one\nTensor of activations per table specified in the model.",
"attributes": [
{
"name": "num_tables",
"type": "int64",
"description": "The number of output activation tensors. If feature descriptor is\npresent in the tpu embedding config, it is equal to the number of features\notherwise equal to number of embedding tables in the model.",
"minimum": 1
},
{
"name": "config",
"type": "string",
"description": "Serialized TPUEmbeddingConfiguration proto."
},
{
"name": "embedding_partitions",
"type": "string",
"description": "Serialized EmbeddingPartitionsProto proto."
},
{
"name": "hbm_buffers_config",
"type": "string",
"description": "Serialized HbmBuffersConfig proto."
},
{
"name": "tpu_topology",
"type": "string",
"description": "Serialized TpuTopologyArgsProto proto."
}
],
"inputs": [
{
"name": "deduplication_data",
"description": "A Tensor with type=DT_VARIANT containing the deduplication\ndata. The tensor is an XLA nested tuple containing N elements (where N is\nthe ratio of the number of embedding to tensor cores per TPU chip). Each\nelement of the nested tuple is a tuple of rank 1 tensors. Each tensor either\ncontains indices (DT_UINT32) for embedding lookup on the TensorCore or\nweights (DT_FLOAT) to apply to the output of the embedding lookup operation.",
"type": 21
}
],
"outputs": [
{
"name": "outputs",
"description": "A TensorList of embedding activations containing one Tensor per\nembedding table in the model.",
"numberAttr": "num_tables",
"type": 1
}
]
},
{
"name": "XlaRecvTPUEmbeddingDeduplicationData",
"summary": "Receives deduplication data (indices and weights) from the embedding core.",
Expand All @@ -64232,6 +64376,39 @@
}
]
},
{
"name": "XlaRecvTPUEmbeddingDeduplicationDataV2",
"summary": "Receives deduplication data (indices and weights) from the embedding core.",
"description": "The deduplication data is a Tensor with type=DT_VARIANT. The tensor itself is an\nXLA nested tuple containing N elements (where N is the ratio of the number of\nembedding to tensor cores per TPU chip). Each element of the nested tuple is a\ntuple of rank 1 tensors. Each tensor either contains indices (DT_UINT32) for\nembedding lookup on the TensorCore or weights (DT_FLOAT) to apply to the output\nof the embedding lookup operation.",
"attributes": [
{
"name": "config",
"type": "string",
"description": "Serialized TPUEmbeddingConfiguration proto."
},
{
"name": "embedding_partitions",
"type": "string",
"description": "Serialized EmbeddingPartitionsProto proto."
},
{
"name": "hbm_buffers_config",
"type": "string",
"description": "Serialized HbmBuffersConfig proto."
},
{
"name": "tpu_topology",
"type": "string",
"description": "Serialized TpuTopologyArgsProto proto."
}
],
"outputs": [
{
"name": "output",
"type": 21
}
]
},
{
"name": "XlaSendTPUEmbeddingGradients",
"summary": "An op that performs gradient updates of embedding tables.",
Expand Down Expand Up @@ -64276,6 +64453,65 @@
}
]
},
{
"name": "XlaSendTPUEmbeddingGradientsV2",
"summary": "An op that performs gradient updates of embedding tables.",
"description": "The gradients argument is a TensorList having the same length and shapes as the\nreturn value of XlaRecvTPUEmbeddingActivations, but contains gradients of the\nmodel's loss with respect to the embedding activations. The embedding tables are\nupdated from these gradients via the optimizer specified in the\nTPUEmbeddingConfiguration proto given to tpu.initialize_system.",
"attributes": [
{
"name": "NumTables",
"type": "int64",
"description": "number of tables",
"minimum": 1
},
{
"name": "NumLearningRateTags",
"type": "int64",
"description": "number of learning rate tags",
"minimum": 0,
"default": 0
},
{
"name": "config",
"type": "string",
"description": "Serialized TPUEmbeddingConfiguration proto."
},
{
"name": "embedding_partitions",
"type": "string",
"description": "Serialized EmbeddingPartitionsProto proto."
},
{
"name": "hbm_buffers_config",
"type": "string",
"description": "Serialized HbmBuffersConfig proto."
},
{
"name": "tpu_topology",
"type": "string",
"description": "Serialized TpuTopologyArgsProto proto."
}
],
"inputs": [
{
"name": "gradients",
"description": "A TensorList of gradients with which to update embedding tables.",
"numberAttr": "NumTables",
"type": 1
},
{
"name": "learning_rates",
"description": "A TensorList of learning rates used for updating the embedding\ntables via the optimizer. The length of the TensorList must be equal to the\nnumber of dynamic learning rate tags specified in the\nTPUEmbeddingConfiguration proto.",
"numberAttr": "NumLearningRateTags",
"type": 1
},
{
"name": "deduplication_data",
"description": "A Tensor with type=DT_VARIANT containing the deduplication\ndata. The tensor is an XLA nested tuple containing N elements (where N is\nthe ratio of the number of embedding to tensor cores per TPU chip). Each\nelement of the nested tuple is a tuple of rank 1 tensors. Each tensor either\ncontains indices (DT_UINT32) for embedding lookup on the TensorCore or\nweights (DT_FLOAT) to apply to the output of the embedding lookup operation.",
"type": 21
}
]
},
{
"name": "XlaSendToHost",
"summary": "An op to send a tensor to the host.",
Expand Down

0 comments on commit 26e138d

Please sign in to comment.