1syntax = "proto3"; 2 3package tensorflow.tpu; 4 5// In the comments here, "layout" refers to the top-level EmbeddingOutputLayout 6// proto contained in the TPUEmbeddingConfiguration. 7 8// The embedding output consists of a list of tensors, each specified by an 9// EmbeddingOutputTensor proto within the EmbeddingOutputLayout (the "output" 10// field). Each table and feature lookup is then placed into some number of 11// particular positions within some output tensor (identified by "tensor_index" 12// within OutputLocation). The tree of table lookups, feature lookups, and 13// output locations is specified by the 14// "table(table_id).feature(feature_id).output_location" repeated fields within 15// EmbeddingOutputLayout. 16 17message TPUEmbeddingOutputLayout { 18 // Location of one copy of the feature's data. 19 message OutputLocation { 20 // Which output tensor this copy of the feature will go into. Must be 21 // between 0 and layout.output_size(). 22 int32 tensor_index = 1; 23 24 // Offset in dimension 0 for this feature copy. Must be between 0 and 25 // layout.output(tensor_index).dim0_size_per_sample(). 26 int32 dim0_offset = 2; 27 28 // Offset in dimension 1 for this feature copy. Must be between 0 and 29 // layout.output(tensor_index).dim1_size() - table width; repeated or 30 // partially/fully overlapping values are allowed and results in the same 31 // range will be summed (with the gradients replicated in the backward 32 // pass). 33 int32 dim1_offset = 3; 34 } 35 36 // Description of the output placement for one feature. 37 message FeatureDescriptor { 38 // Typically, only one copy of each feature is used, but multiple are 39 // allowed and the same data will be copied to all of them (with the 40 // gradients summed in the backward pass). 41 repeated OutputLocation output_location = 1; 42 } 43 44 // Description of the output placement for features of one table. 45 message TableDescriptor { 46 // Output locations for each feature loaded from this table. 47 repeated FeatureDescriptor feature = 1; 48 } 49 // Output locations for each feature of each table. 50 repeated TableDescriptor table = 1; 51 52 // Data layout and shape computation information for a single output tensor. 53 // Any unused locations in the tensor will be filled with zeros, and 54 // corresponding gradients will be ignored. 55 56 // Size and layout information for 2-D tensors. 57 message TwoDOutputTensor { 58 // Multiplier for output dimension 0 size; used to match legacy format that 59 // stacks features within a sample in dimension 0. 60 int32 dim0_size_per_sample = 2; 61 62 // The size (in dimension 1) of this output tensor. 63 int32 dim1_size = 1; 64 } 65 66 // Format information for a single output tensor. 67 message EmbeddingOutputTensor { 68 oneof output_format { 69 TwoDOutputTensor two_d = 4; 70 } 71 } 72 73 // Shape and layout information for each tensor. 74 repeated EmbeddingOutputTensor output = 2; 75} 76