From e802b0498f3ccffda31a33a5ef591514c25b3cc9 Mon Sep 17 00:00:00 2001
From: RandySheriffH <48490400+RandySheriffH@users.noreply.github.com>
Date: Wed, 5 Aug 2020 14:11:56 -0700
Subject: [PATCH] EnrichPyOpUT (#4681)

* cancel night build on pyop

* enrich PyOp UTs

* init script only once

* remove space

* update models

* Show usage of kwargs in doc
---
 docs/PyOp.md                                  |  5 +-
 onnxruntime/test/shared_lib/test_inference.cc | 54 ++++++++++++++++++-
 onnxruntime/test/testdata/pyop_2.onnx         | 31 +++++++++++
 onnxruntime/test/testdata/pyop_3.onnx         | 30 +++++++++++
 4 files changed, 115 insertions(+), 5 deletions(-)
 create mode 100644 onnxruntime/test/testdata/pyop_2.onnx
 create mode 100644 onnxruntime/test/testdata/pyop_3.onnx

diff --git a/docs/PyOp.md b/docs/PyOp.md
index 2380d48eab..0764ff8b3d 100644
--- a/docs/PyOp.md
+++ b/docs/PyOp.md
@@ -57,9 +57,8 @@ class Multi_1:
         ret = S + P
         return ret + self.W1, ret + self.W2, ret + self.W3
 class Multi_2:
-    def compute(self, H, N, E):
-        r1, r2 = H + N, N + E
-        return r1, r2
+    def compute(self, *kwargs):
+        return sum(kwargs[0:-1]), sum(kwargs[1:])
 ```
 ### Step 4
 Copy mymodule.py into Python sys.path, then run the model with onnxruntime python API. On Windows, please set PYTHONHOME beforehand. It should point to directory where the python is installed, such as C:\Python37 or C:\ProgramData\Anaconda3\envs\myconda1 if it is in conda.
diff --git a/onnxruntime/test/shared_lib/test_inference.cc b/onnxruntime/test/shared_lib/test_inference.cc
index 0dbbe219d4..e21fc9675c 100644
--- a/onnxruntime/test/shared_lib/test_inference.cc
+++ b/onnxruntime/test/shared_lib/test_inference.cc
@@ -144,6 +144,8 @@ static constexpr PATH_TYPE MODEL_WITH_CUSTOM_MODEL_METADATA = TSTR("testdata/mod
 
 #ifdef ENABLE_LANGUAGE_INTEROP_OPS
 static constexpr PATH_TYPE PYOP_FLOAT_MODEL_URI = TSTR("testdata/pyop_1.onnx");
+static constexpr PATH_TYPE PYOP_MULTI_MODEL_URI = TSTR("testdata/pyop_2.onnx");
+static constexpr PATH_TYPE PYOP_KWARG_MODEL_URI = TSTR("testdata/pyop_3.onnx");
 #endif
 
 class CApiTestWithProvider : public testing::Test, public ::testing::WithParamInterface<int> {
@@ -359,8 +361,10 @@ TEST(CApiTest, test_custom_op_library) {
 }
 
 #if defined(ENABLE_LANGUAGE_INTEROP_OPS)
-TEST(CApiTest, test_pyop) {
-  std::cout << "Test model with pyop" << std::endl;
+std::once_flag my_module_flag;
+
+void PrepareModule()
+{
   std::ofstream module("mymodule.py");
   module << "class MyKernel:" << std::endl;
   module << "\t"
@@ -371,7 +375,29 @@ TEST(CApiTest, test_pyop) {
          << "def compute(self,x):" << std::endl;
   module << "\t\t"
          << "return x*2" << std::endl;
+  module << "class MyKernel_2:" << std::endl;
+  module << "\t"
+         << "def __init__(self,A,B):" << std::endl;
+  module << "\t\t"
+         << "self.a,self.b = A,B" << std::endl;
+  module << "\t"
+         << "def compute(self,x):" << std::endl;
+  module << "\t\t"
+         << "return x*4" << std::endl;
+  module << "class MyKernel_3:" << std::endl;
+  module << "\t"
+         << "def __init__(self,A,B):" << std::endl;
+  module << "\t\t"
+         << "self.a,self.b = A,B" << std::endl;
+  module << "\t"
+         << "def compute(self,*kwargs):" << std::endl;
+  module << "\t\t"
+         << "return kwargs[0]*5" << std::endl;
   module.close();
+}
+
+TEST(CApiTest, test_pyop) {
+  std::call_once(my_module_flag, PrepareModule);
   std::vector<Input> inputs(1);
   Input& input = inputs[0];
   input.name = "X";
@@ -381,6 +407,30 @@ TEST(CApiTest, test_pyop) {
   std::vector<float> expected_values_y = {2.0f, 4.0f, 6.0f, 8.0f};
   TestInference<PATH_TYPE, float>(*ort_env, PYOP_FLOAT_MODEL_URI, inputs, "Y", expected_dims_y, expected_values_y, 0, nullptr, nullptr);
 }
+
+TEST(CApiTest, test_pyop_multi) {
+  std::call_once(my_module_flag, PrepareModule);
+  std::vector<Input> inputs(1);
+  Input& input = inputs[0];
+  input.name = "X";
+  input.dims = {2, 2};
+  input.values = {1.0f, 2.0f, 3.0f, 4.0f};
+  std::vector<int64_t> expected_dims_y = {2, 2};
+  std::vector<float> expected_values_y = {8.0f, 16.0f, 24.0f, 32.0f};
+  TestInference<PATH_TYPE, float>(*ort_env, PYOP_MULTI_MODEL_URI, inputs, "Z", expected_dims_y, expected_values_y, 0, nullptr, nullptr);
+}
+
+TEST(CApiTest, test_pyop_kwarg) {
+  std::call_once(my_module_flag, PrepareModule);
+  std::vector<Input> inputs(1);
+  Input& input = inputs[0];
+  input.name = "X";
+  input.dims = {2, 2};
+  input.values = {1.0f, 2.0f, 3.0f, 4.0f};
+  std::vector<int64_t> expected_dims_y = {2, 2};
+  std::vector<float> expected_values_y = {25.0f, 50.0f, 75.0f, 100.0f};
+  TestInference<PATH_TYPE, float>(*ort_env, PYOP_KWARG_MODEL_URI, inputs, "Z", expected_dims_y, expected_values_y, 0, nullptr, nullptr);
+}
 #endif
 
 #ifdef ORT_RUN_EXTERNAL_ONNX_TESTS
diff --git a/onnxruntime/test/testdata/pyop_2.onnx b/onnxruntime/test/testdata/pyop_2.onnx
new file mode 100644
index 0000000000..51d4e91133
--- /dev/null
+++ b/onnxruntime/test/testdata/pyop_2.onnx
@@ -0,0 +1,31 @@
+
+pyop_model:�
+�

+
X
Y"PyOp*	
+
A"
1�
*	
+
B"
2�
*	
+
C"
3�
*
+
+class_name"MyKernel�
*
+input_types@
�
*
+module"mymodule�
*
+output_types@
�
:pyop_multi_1
+�

+
Y
Z"PyOp*	
+
A"
1�
*
+
+
B"11�
*
+
+class_name"
+MyKernel_2�
*
+input_types@
�
*
+module"mymodule�
*
+output_types@
�
:pyop_multi_2multi_pyop_graphZ
+
X
+

+
+b
+
Z
+

+
+B
\ No newline at end of file
diff --git a/onnxruntime/test/testdata/pyop_3.onnx b/onnxruntime/test/testdata/pyop_3.onnx
new file mode 100644
index 0000000000..6b69f58c11
--- /dev/null
+++ b/onnxruntime/test/testdata/pyop_3.onnx
@@ -0,0 +1,30 @@
+
+pyop_model:�
+�

+
X
Y"PyOp*	
+
A"
1�
*	
+
B"
2�
*
+
+class_name"
+MyKernel_3�
*
+input_types@
�
*
+module"mymodule�
*
+output_types@
�
:pyop_kwarg_1
+�

+
Y
Z"PyOp*	
+
A"
2�
*	
+
B"
1�
*
+
+class_name"
+MyKernel_3�
*
+input_types@
�
*
+module"mymodule�
*
+output_types@
�
:pyop_kwarg_2multi_pyop_graphZ
+
X
+

+
+b
+
Z
+

+
+B
\ No newline at end of file