| op { |
| name: "EnqueueTPUEmbeddingSparseBatch" |
| input_arg { |
| name: "sample_indices" |
| type: DT_INT32 |
| number_attr: "N" |
| } |
| input_arg { |
| name: "embedding_indices" |
| type: DT_INT32 |
| number_attr: "N" |
| } |
| input_arg { |
| name: "aggregation_weights" |
| type: DT_FLOAT |
| number_attr: "N" |
| } |
| input_arg { |
| name: "mode_override" |
| type: DT_STRING |
| } |
| attr { |
| name: "N" |
| type: "int" |
| has_minimum: true |
| minimum: 1 |
| } |
| attr { |
| name: "device_ordinal" |
| type: "int" |
| default_value { |
| i: -1 |
| } |
| } |
| attr { |
| name: "combiners" |
| type: "list(string)" |
| default_value { |
| list { |
| } |
| } |
| } |
| is_stateful: true |
| } |
| op { |
| name: "EnqueueTPUEmbeddingSparseBatch" |
| input_arg { |
| name: "sample_indices" |
| type_attr: "T1" |
| number_attr: "N" |
| } |
| input_arg { |
| name: "embedding_indices" |
| type_attr: "T2" |
| number_attr: "N" |
| } |
| input_arg { |
| name: "aggregation_weights" |
| type_attr: "T3" |
| number_attr: "N" |
| } |
| input_arg { |
| name: "mode_override" |
| type: DT_STRING |
| } |
| attr { |
| name: "T1" |
| type: "type" |
| default_value { |
| type: DT_INT32 |
| } |
| allowed_values { |
| list { |
| type: DT_INT32 |
| type: DT_INT64 |
| } |
| } |
| } |
| attr { |
| name: "T2" |
| type: "type" |
| default_value { |
| type: DT_INT32 |
| } |
| allowed_values { |
| list { |
| type: DT_INT32 |
| type: DT_INT64 |
| } |
| } |
| } |
| attr { |
| name: "T3" |
| type: "type" |
| default_value { |
| type: DT_FLOAT |
| } |
| allowed_values { |
| list { |
| type: DT_FLOAT |
| type: DT_DOUBLE |
| } |
| } |
| } |
| attr { |
| name: "N" |
| type: "int" |
| has_minimum: true |
| minimum: 1 |
| } |
| attr { |
| name: "device_ordinal" |
| type: "int" |
| default_value { |
| i: -1 |
| } |
| } |
| attr { |
| name: "combiners" |
| type: "list(string)" |
| default_value { |
| list { |
| } |
| } |
| } |
| is_stateful: true |
| } |