From 024476556f62a6f0f2d2ff83ba4d4be423061cc8 Mon Sep 17 00:00:00 2001
From: Josh Fromm <jwfromm@uw.edu>
Date: Thu, 3 Dec 2020 19:00:16 -0800
Subject: [PATCH 1/2] Add support for Size op in Onnx frontend.

---
 python/tvm/relay/frontend/onnx.py          |  1 +
 tests/python/frontend/onnx/test_forward.py | 31 ++++++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py
index 7ebad7297471..d65f5676fb33 100644
--- a/python/tvm/relay/frontend/onnx.py
+++ b/python/tvm/relay/frontend/onnx.py
@@ -2374,6 +2374,7 @@ def _get_convert_map(opset):
         "Gather": Gather.get_converter(opset),
         "GatherElements": GatherElements.get_converter(opset),
         "GatherND": GatherND.get_converter(opset),
+        "Size": AttrCvt("ndarray_size", extras={"dtype": "int64"}),
         "Scatter": Scatter.get_converter(opset),
         "ScatterElements": Scatter.get_converter(opset),
         "Squeeze": AttrCvt("squeeze", {"axes": "axis"}),
diff --git a/tests/python/frontend/onnx/test_forward.py b/tests/python/frontend/onnx/test_forward.py
index b84e55ac800c..1b401dfec967 100644
--- a/tests/python/frontend/onnx/test_forward.py
+++ b/tests/python/frontend/onnx/test_forward.py
@@ -3888,6 +3888,36 @@ def test_if():
             tvm.testing.assert_allclose(correct_out[i], tvm_out[i], rtol=1e-05, atol=1e-05)
 
 
+@tvm.testing.uses_gpu
+def test_size():
+    def verify_size(indata):
+        node = helper.make_node(
+            "Size",
+            inputs=["X"],
+            outputs=["Y"],
+        )
+
+        graph = helper.make_graph(
+            [node],
+            "size_test",
+            inputs=[helper.make_tensor_value_info("X", TensorProto.INT64, list(indata.shape))],
+            outputs=[helper.make_tensor_value_info("Y", TensorProto.INT64, [])],
+        )
+
+        model = helper.make_model(graph, producer_name="size_test")
+
+        for target, _ in tvm.testing.enabled_targets():
+            verify_with_ort_with_inputs(
+                model, [indata], targets=[target], dtype="int64", use_vm=True, opset=11
+            )
+
+    input_data = np.array([[1, 0], [1, 1]], dtype=np.int64)
+    verify_size(input_data)
+
+    input_data = np.array([[3, 0, 0], [0, 4, 0], [5, 6, 0]], dtype=np.int64)
+    verify_size(input_data)
+
+
 if __name__ == "__main__":
     test_flatten()
     test_reshape()
@@ -3964,3 +3994,4 @@ def test_if():
     test_roi_align()
     test_range()
     test_loop()
+    test_size()

From e86bc1816972f1900aded4b0cd6eb82f83310a46 Mon Sep 17 00:00:00 2001
From: Josh Fromm <jwfromm@uw.edu>
Date: Thu, 3 Dec 2020 19:33:43 -0800
Subject: [PATCH 2/2] Simplify target testing.

---
 tests/python/frontend/onnx/test_forward.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/tests/python/frontend/onnx/test_forward.py b/tests/python/frontend/onnx/test_forward.py
index 1b401dfec967..3ddc80af3a32 100644
--- a/tests/python/frontend/onnx/test_forward.py
+++ b/tests/python/frontend/onnx/test_forward.py
@@ -3906,10 +3906,7 @@ def verify_size(indata):
 
         model = helper.make_model(graph, producer_name="size_test")
 
-        for target, _ in tvm.testing.enabled_targets():
-            verify_with_ort_with_inputs(
-                model, [indata], targets=[target], dtype="int64", use_vm=True, opset=11
-            )
+        verify_with_ort_with_inputs(model, [indata], dtype="int64", use_vm=True, opset=11)
 
     input_data = np.array([[1, 0], [1, 1]], dtype=np.int64)
     verify_size(input_data)