Adding TorchEmbedding contrib op (#7136)

* Adding TorchEmbedding contrib op

* Update contrib_defs.cc

* Shape fix

* Update shape_inference_test_helper.h

* Fix typo

* Fix test

* Fix for test code

* Merge

* Fix CI

* Fix for CI

* Fix CI no-contrib

onnxruntime/core/graph/contrib_ops/contrib_defs.cc

@@ -2290,6 +2290,83 @@ It's an extension of Gelu. It takes the sum of input A and bias input B as the i
         }
       });
 
+  static const char* TorchEmbedding_ver1_doc = R"DOC(
+      Based on Torch operator Embedding, creates a lookup table of embedding vectors of fixed size,
+       for a dictionary of fixed size.
+      )DOC";
+
+  ONNX_CONTRIB_OPERATOR_SCHEMA(TorchEmbedding)
+      .SetDomain(kMSDomain)
+      .SinceVersion(1)
+      .SetDoc(TorchEmbedding_ver1_doc)
+      .Input(
+          0,
+          "weight",
+          "The embedding matrix of size N x M. 'N' is equal to the maximum possible index + 1, and 'M' is "
+          "equal to the embedding size",
+          "T")
+      .Input(
+          1,
+          "indices",
+          "Long tensor containing the indices to extract from embedding matrix.",
+          "tensor(int64)")
+      .Input(
+          2,
+          "padding_idx",
+          "A 0-D scalar tensor. If specified, the entries at `padding_idx` do not contribute to the gradient; "
+          "therefore, the embedding vector at `padding_idx` is not updated during training, "
+          "i.e. it remains as a fixed pad.",
+          "tensor(int64)",
+          OpSchema::Optional)
+      .Input(
+          3,
+          "scale_grad_by_freq",
+          "A 0-D bool tensor. If given, this will scale gradients by the inverse of frequency of "
+          "the indices (words) in the mini-batch. Default  is ``False``",
+          "tensor(bool)",
+          OpSchema::Optional)
+      .Output(
+          0,
+          "Y",
+          "Output tensor of the same type as the input tensor. Shape of the output is * x M, where '*' is the shape of "
+          "input indices, and 'M' is the embedding size.",
+          "T")
+      .TypeConstraint(
+          "T",
+          {"tensor(float16)",
+           "tensor(float)",
+           "tensor(double)",
+           "tensor(bfloat16)",
+           "tensor(uint8)",
+           "tensor(uint16)",
+           "tensor(uint32)",
+           "tensor(uint64)",
+           "tensor(int8)",
+           "tensor(int16)",
+           "tensor(int32)",
+           "tensor(int64)"},
+          "Constrain input and output types to all numeric tensors.")
+      .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
+        using namespace ONNX_NAMESPACE;
+        propagateElemTypeFromInputToOutput(ctx, 0, 0);
+
+        TensorShapeProto outputs_shape;
+        Dim input_dim_i;
+
+        if (hasInputShape(ctx, 1)) {
+          auto& input_shape = getInputShape(ctx, 1);
+          for (int32_t i = 0; i < input_shape.dim_size(); i++) {
+            input_dim_i = input_shape.dim(i);
+            *outputs_shape.add_dim() = input_dim_i;
+          }
+        }
+
+        Dim embedding_dim;
+        unifyInputDim(ctx, 0, 1, embedding_dim);
+        *outputs_shape.add_dim() = embedding_dim;
+        updateOutputShape(ctx, 0, outputs_shape);
+      });
+
   static const char* Trilu_ver1_doc = R"DOC(
       Returns the upper or lower triangular part of a 2-D matrix, or batches of 2-D matrices. If the attribute "upper" is set to true,
       the upper triangular matrix is retained. Lower triangular matrix is retained otherwise. Default value for upper is true.

onnxruntime/test/contrib_ops/torch_embedding_shape_inference_test.cc

@@ -0,0 +1,109 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "test/providers/provider_test_utils.h"
+#include "onnx/shape_inference/implementation.h"
+#include "onnx/checker.h"
+#include "test/providers/cpu/tensor/shape_inference_test_helper.h"
+
+namespace onnxruntime {
+namespace test {
+
+TEST(ShapeInferenceTests, embedding_float) {
+  std::vector<int64_t> shape = {2, 4};
+  ONNX_NAMESPACE::ValueInfoProto input;
+  CreateValueInfo(input, "indices", ONNX_NAMESPACE::TensorProto_DataType_INT64, shape);
+  std::vector<int64_t> shape_w = {10, 3};
+  ONNX_NAMESPACE::ValueInfoProto weight;
+  CreateValueInfo(weight, "weight", ONNX_NAMESPACE::TensorProto_DataType_FLOAT, shape_w);
+  std::vector<ONNX_NAMESPACE::ValueInfoProto> inputs = {weight, input};
+
+  std::vector<ONNX_NAMESPACE::AttributeProto> attributes = {};
+
+  std::vector<int64_t> shape_y = {2, 4, 3};
+  ONNX_NAMESPACE::ValueInfoProto output;
+  CreateValueInfo(output, "Y", ONNX_NAMESPACE::TensorProto_DataType_FLOAT, shape_y);
+
+  TestShapeInference("TorchEmbedding", kMSDomain, 1, 6, inputs, attributes, output);
+}
+
+TEST(ShapeInferenceTests, embedding_zero_dim_int) {
+  std::vector<int64_t> shape = {0, 3};
+  ONNX_NAMESPACE::ValueInfoProto input;
+  CreateValueInfo(input, "X", ONNX_NAMESPACE::TensorProto_DataType_INT64, shape);
+  std::vector<int64_t> shape_w = {10, 4};
+  ONNX_NAMESPACE::ValueInfoProto weight;
+  CreateValueInfo(weight, "W", ONNX_NAMESPACE::TensorProto_DataType_INT32, shape_w);
+  std::vector<ONNX_NAMESPACE::ValueInfoProto> inputs = {weight, input};
+
+  std::vector<ONNX_NAMESPACE::AttributeProto> attributes = {};
+
+  std::vector<int64_t> shape_y = {0, 3, 4};
+  ONNX_NAMESPACE::ValueInfoProto output;
+  CreateValueInfo(output, "Y", ONNX_NAMESPACE::TensorProto_DataType_INT32, shape_y);
+
+  TestShapeInference("TorchEmbedding", kMSDomain, 1, 6, inputs, attributes, output);
+}
+
+TEST(ShapeInferenceTests, embedding_long) {
+  std::vector<int64_t> shape = {2, 4};
+  ONNX_NAMESPACE::ValueInfoProto input;
+  CreateValueInfo(input, "X", ONNX_NAMESPACE::TensorProto_DataType_INT64, shape);
+  std::vector<int64_t> shape_w = {10, 3};
+  ONNX_NAMESPACE::ValueInfoProto weight;
+  CreateValueInfo(weight, "W", ONNX_NAMESPACE::TensorProto_DataType_INT64, shape_w);
+  std::vector<ONNX_NAMESPACE::ValueInfoProto> inputs = {weight, input};
+
+  std::vector<ONNX_NAMESPACE::AttributeProto> attributes = {};
+
+  std::vector<int64_t> shape_y = {2, 4, 3};
+  ONNX_NAMESPACE::ValueInfoProto output;
+  CreateValueInfo(output, "Y", ONNX_NAMESPACE::TensorProto_DataType_INT64, shape_y);
+
+  TestShapeInference("TorchEmbedding", kMSDomain, 1, 6, inputs, attributes, output);
+}
+
+TEST(ShapeInferenceTests, embedding_with_padding) {
+  std::vector<int64_t> shape = {2, 4};
+  ONNX_NAMESPACE::ValueInfoProto input;
+  CreateValueInfo(input, "indices", ONNX_NAMESPACE::TensorProto_DataType_INT64, shape);
+  std::vector<int64_t> shape_w = {10, 3};
+  ONNX_NAMESPACE::ValueInfoProto weight;
+  CreateValueInfo(weight, "weight", ONNX_NAMESPACE::TensorProto_DataType_INT64, shape_w);
+  ONNX_NAMESPACE::ValueInfoProto padding_idx;
+  CreateValueInfo(padding_idx, "Pad", ONNX_NAMESPACE::TensorProto_DataType_INT64, {});
+  std::vector<ONNX_NAMESPACE::ValueInfoProto> inputs = {weight, input, padding_idx};
+
+  std::vector<ONNX_NAMESPACE::AttributeProto> attributes = {};
+
+  std::vector<int64_t> shape_y = {2, 4, 3};
+  ONNX_NAMESPACE::ValueInfoProto output;
+  CreateValueInfo(output, "Y", ONNX_NAMESPACE::TensorProto_DataType_INT64, shape_y);
+
+  TestShapeInference("TorchEmbedding", kMSDomain, 1, 6, inputs, attributes, output);
+}
+
+TEST(ShapeInferenceTests, embedding_with_scale_grad) {
+  std::vector<int64_t> shape = {2, 8, 4};
+  ONNX_NAMESPACE::ValueInfoProto input;
+  CreateValueInfo(input, "indices", ONNX_NAMESPACE::TensorProto_DataType_INT64, shape);
+  std::vector<int64_t> shape_w = {10, 3};
+  ONNX_NAMESPACE::ValueInfoProto weight;
+  CreateValueInfo(weight, "weight", ONNX_NAMESPACE::TensorProto_DataType_INT64, shape_w);
+  ONNX_NAMESPACE::ValueInfoProto padding_idx;
+  CreateValueInfo(padding_idx, "Pad", ONNX_NAMESPACE::TensorProto_DataType_INT64, {});
+  ONNX_NAMESPACE::ValueInfoProto scale_grad;
+  CreateValueInfo(scale_grad, "Scale", ONNX_NAMESPACE::TensorProto_DataType_BOOL, {});
+  std::vector<ONNX_NAMESPACE::ValueInfoProto> inputs = {weight, input, padding_idx, scale_grad};
+
+  std::vector<ONNX_NAMESPACE::AttributeProto> attributes = {};
+
+  std::vector<int64_t> shape_y = {2, 8, 4, 3};
+  ONNX_NAMESPACE::ValueInfoProto output;
+  CreateValueInfo(output, "Y", ONNX_NAMESPACE::TensorProto_DataType_INT64, shape_y);
+
+  TestShapeInference("TorchEmbedding", kMSDomain, 1, 6, inputs, attributes, output);
+}
+
+}  // namespace test
+}  // namespace onnxruntime

onnxruntime/test/providers/cpu/tensor/shape_inference_test_helper.h

@@ -9,9 +9,9 @@
 namespace onnxruntime {
 namespace test {
 
-auto schema_registry = ONNX_NAMESPACE::OpSchemaRegistry::Instance();
+static auto schema_registry = ONNX_NAMESPACE::OpSchemaRegistry::Instance();
 
-void CheckShapeEquality(ONNX_NAMESPACE::TensorShapeProto* shape1, ONNX_NAMESPACE::TensorShapeProto* shape2) {
+inline void CheckShapeEquality(ONNX_NAMESPACE::TensorShapeProto* shape1, ONNX_NAMESPACE::TensorShapeProto* shape2) {
   EXPECT_NE(shape1, nullptr);
   EXPECT_NE(shape2, nullptr);
   if ((shape1 != nullptr) && (shape2 != nullptr)) {
@@ -66,29 +66,31 @@ inline void TestShapeInference(const std::string& op_type,
   // Set model graph
   ONNX_NAMESPACE::GraphProto* graph = model.mutable_graph();
   graph->set_name("test-op");
-  graph->add_value_info();
 
   // Set add operator node to graph
-  auto& node = *graph->add_node();
-  node.set_op_type(op_type);
-  node.set_domain(op_domain);
-  node.set_name("test_node");
+  auto node = graph->add_node();
+  node->set_op_type(op_type);
+  node->set_domain(op_domain);
+  node->set_name("test_node");
 
   // Add node inputs and graph inputs
-  for (auto const& n_ : inputs) {
-    node.add_input(n_.name());
-    *graph->add_input() = n_;
-  }
+	for (auto const& n_ : inputs) {
+	  node->add_input(n_.name());
+	  auto in = graph->add_input();
+	  *in = n_;
+	  auto v_ = graph->add_value_info();
+	  *v_ = n_;
+	}
 
   // Add node attributes
   for (auto const& attr : attributes) {
-    node.add_attribute()->CopyFrom(attr);
+    node->add_attribute()->CopyFrom(attr);
   }
 
-  node.add_output("Output");
+  node->add_output("Output");
 
+  ONNX_NAMESPACE::shape_inference::InferShapes(model, true, schema_registry);
   ONNX_NAMESPACE::checker::check_model(model);
-  ONNX_NAMESPACE::shape_inference::InferShapes(model, false, schema_registry);
 
   auto inferredGraph = model.graph();
   int index = static_cast<int>(inputs.size());  // index for value_info of output

onnxruntime/test/python/contrib_ops/onnx_test_torch_embedding.py

@@ -0,0 +1,87 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+#
+# Test reference implementation and model for ONNX Runtime conrtib op torch_embedding
+
+import onnx
+import unittest
+import numpy as np
+from onnx_contrib_ops_helper import expect
+
+
+def torch_embedding_reference_implementation(weight, indices, padding_idx=None, scale=False):
+    return np.take(weight, indices, axis=0)
+
+
+class ONNXReferenceImplementationTest(unittest.TestCase):
+    def test_torch_embedding(self):
+        node = onnx.helper.make_node(
+            'TorchEmbedding',
+            inputs=['w', 'x'],
+            outputs=['y'],
+            domain="com.microsoft",
+        )
+
+        x = np.random.randn(2, 4).astype(np.int64)
+        w = np.random.randn(10, 3).astype(np.float32)
+        y = torch_embedding_reference_implementation(w, x)
+        expect(node, inputs=[w, x], outputs=[y], name='test_torch_embedding')
+
+    def test_torch_embedding_long(self):
+        node = onnx.helper.make_node(
+            'TorchEmbedding',
+            inputs=['w', 'x'],
+            outputs=['y'],
+            domain="com.microsoft",
+        )
+
+        x = np.random.randn(2, 4).astype(np.int64)
+        w = np.random.randn(10, 3).astype(np.int64)
+        y = torch_embedding_reference_implementation(w, x)
+        expect(node, inputs=[w, x], outputs=[y], name='test_torch_embedding_long')
+
+    def test_torch_embedding_zero_dim(self):
+        node = onnx.helper.make_node(
+            'TorchEmbedding',
+            inputs=['w', 'x'],
+            outputs=['y'],
+            domain="com.microsoft",
+        )
+
+        x = np.random.randn(0, 4).astype(np.int64)
+        w = np.random.randn(10, 3).astype(np.float32)
+        y = torch_embedding_reference_implementation(w, x)
+        expect(node, inputs=[w, x], outputs=[y], name='test_torch_embedding_zero_dim')
+
+    def test_torch_embedding_padding_idx(self):
+        node = onnx.helper.make_node(
+            'TorchEmbedding',
+            inputs=['w', 'x', 'padding_idx'],
+            outputs=['y'],
+            domain="com.microsoft",
+        )
+
+        x = np.random.randn(3, 4).astype(np.int64)
+        w = np.random.randn(10, 3).astype(np.float32)
+        padding_idx = np.random.randint(3, size=1).astype(np.int64)
+        y = torch_embedding_reference_implementation(w, x, padding_idx)
+        expect(node, inputs=[w, x, padding_idx], outputs=[y], name='test_torch_embedding_padding_idx')
+
+    def test_torch_embedding_scale_grad_by_freq(self):
+        node = onnx.helper.make_node(
+            'TorchEmbedding',
+            inputs=['w', 'x', 'padding_idx', 'scale'],
+            outputs=['y'],
+            domain="com.microsoft",
+        )
+
+        x = np.random.randn(3, 4).astype(np.int64)
+        w = np.random.randn(10, 3).astype(np.float32)
+        padding_idx = np.random.randint(3, size=1).astype(np.int64)
+        scale = np.array([1]).astype(np.bool)
+        y = torch_embedding_reference_implementation(w, x, padding_idx, scale)
+        expect(node, inputs=[w, x, padding_idx, scale], outputs=[y], name='test_torch_embedding_scale_grad_by_freq')
+
+
+if __name__ == '__main__':
+    unittest.main(module=__name__, buffer=True)

onnxruntime/test/python/testdata/test_torch_embedding/test_data_set_0/input_0.pb

@@ -0,0 +1,2 @@
+
+
B
wJxֻור>YWֻ¾RCp?_Rִ?„:׃¾שB> ³¾±<C2BE>L¿ְ`¸¾‡"<22>¿׀<C2BF>ז>8½>a‡?2j¿÷‘י?­­<C2AD>¿ּ?*ײױ><3E>X?$?>ײ¿	<09>>©<>h?<02>E¿hI±¾1?ˆqG>M ¿<>&<26><M»„¿

onnxruntime/test/python/testdata/test_torch_embedding/test_data_set_0/output_0.pb

@@ -0,0 +1 @@
+
B
yJ`ֻור>YWֻ¾RCp?ֻור>YWֻ¾RCp?M ¿<>&<26><M»„¿ֻור>YWֻ¾RCp?ֻור>YWֻ¾RCp?ֻור>YWֻ¾RCp?_Rִ?„:׃¾שB>M ¿<>&<26><M»„¿

onnxruntime/test/python/testdata/test_torch_embedding_padding_idx/test_data_set_0/input_0.pb

@@ -0,0 +1,2 @@
+
+
B
wJx+yEї?1g?Южa>°Ћџї¤Є=вО@їф+=НlН>,‘ї!ЂµѕcѕыqLї·"?”@ы–@ЁЮyѕspЉї[ЫЊ?f0ьїќЬЧ>ъьYїdЛ>с1‹?М†ЉїУs‹ѕНҐ*їТ8©ѕiђ\ѕ‚\©>¤И`?

onnxruntime/test/python/testdata/test_torch_embedding_padding_idx/test_data_set_0/output_0.pb

@@ -0,0 +1 @@
+
B
yJђ
+yEї?1g?Южa>°Ћџї¤Є=вО@ї+yEї?1g?Южa>+yEї?1g?Южa>+yEї?1g?Южa>+yEї?1g?Южa>+yEї?1g?Южa>iђ\ѕ‚\©>¤И`?iђ\ѕ‚\©>¤И`?+yEї?1g?Южa>°Ћџї¤Є=вО@їф+=НlН>,‘ї

onnxruntime/test/python/testdata/test_torch_embedding_scale_grad_by_freq/test_data_set_0/input_0.pb

@@ -0,0 +1,2 @@
+
+
B
wJx‡Ώp"ΐ¤w>—<>-Ύ³[?K•LΎ“=ΏK„?¶[<¬…ΏΗψ‘?NϋΎΣχEΎΧ„>ή<>?ΓΏD=JS?Ρ.€ΎμΘ?®}‚?}
·>…-Ώ5³μΏu.Ό?v/ο?ίζzΏHξκΎ|¨?Υόο?ο£]Ώ

onnxruntime/test/python/testdata/test_torch_embedding_scale_grad_by_freq/test_data_set_0/input_3.pb

@@ -0,0 +1 @@
+
	BscaleJ

onnxruntime/test/python/testdata/test_torch_embedding_scale_grad_by_freq/test_data_set_0/output_0.pb

@@ -0,0 +1 @@
+
B
yJ<79>
‡¿p"À¤w>‡¿p"À¤w>‡¿p"À¤w>‡¿p"À¤w>‡¿p"À¤w>‡¿p"À¤w>|¨?Õüï?ï£]¿‡¿p"À¤w>‡¿p"À¤w>—<>-¾³[?K•L¾‡¿p"À¤w>‡¿p"À¤w>

onnxruntime/test/python/testdata/test_torch_embedding_zero_dim/test_data_set_0/input_0.pb

@@ -0,0 +1,2 @@
+
+
B
wJxcżC@7<>–?¬ç•ľěĐ󿮨r>VŰľ|ÔÁ?Ć’1żM?Đ>Ţł<C5A2>ľđ±Ő>7;!żé±GżÇÔ1żb^bż{ňť?¶)Ľ?Öc~>v<>/ľ^a—żĽOg?ý“Bżń%
? Ó6ż&@Lq™>*<2A>6?ëSżw$ś?Ţ	Uż