extend unit test copverage

sdpython · sdpython · Feb 15, 2024 · Feb 2, 2024 · Feb 2, 2024 · Feb 2, 2024
commit 5f37a5997f6ee6755733ee13f3c307d0ddb0e0c3
diff --git a/CHANGELOGS.rst b/CHANGELOGS.rst
@@ -4,6 +4,7 @@ Change Logs
 0.2.0
 +++++
 
+* :pr:`75`: add QuickGelu to ExtendedReferenceEvaluator
 * :pr:`71`: adds tools to compare two onnx graphs
 * :pr:`61`: adds function to plot onnx model as graphs
 * :pr:`60`: supports translation of local functions

diff --git a/_unittests/ut_reference/test_reference_ops.py b/_unittests/ut_reference/test_reference_ops.py
@@ -59,6 +59,25 @@ def test_fused_matmul11(self):
         got = ref.run(None, {"X": a, "Y": a})
         self.assertEqualArray(a.T @ a.T, got[0])
 
+    def test_memcpy(self):
+        model = make_model(
+            make_graph(
+                [
+                    make_node("MemcpyToHost", ["X"], ["Z"]),
+                    make_node("MemcpyFromHost", ["X"], ["Z"]),
+                ],
+                "name",
+                [make_tensor_value_info("X", TensorProto.FLOAT, None)],
+                [make_tensor_value_info("Z", TensorProto.FLOAT, None)],
+            ),
+            opset_imports=[make_opsetid("", 18), make_opsetid("com.microsoft", 1)],
+            ir_version=9,
+        )
+        a = np.arange(4).reshape(-1, 2).astype(np.float32)
+        ref = ExtendedReferenceEvaluator(model)
+        got = ref.run(None, {"X": a})
+        self.assertEqualArray(a, got[0])
+
     def test_quick_gelu(self):
         from onnxruntime import InferenceSession
 

diff --git a/onnx_array_api/reference/evaluator.py b/onnx_array_api/reference/evaluator.py
@@ -8,8 +8,9 @@
 from .ops.op_concat import Concat
 from .ops.op_constant_of_shape import ConstantOfShape
 from .ops.op_fused_matmul import FusedMatMul
-from .ops.op_scatter_elements import ScatterElements
+from .ops.op_memcpy_host import MemcpyFromHost, MemcpyToHost
 from .ops.op_quick_gelu import QuickGelu
+from .ops.op_scatter_elements import ScatterElements
 
 
 logger = getLogger("onnx-array-api-eval")
@@ -36,6 +37,8 @@ class ExtendedReferenceEvaluator(ReferenceEvaluator):
         CastLike_19,
         ConstantOfShape,
         FusedMatMul,
+        MemcpyFromHost,
+        MemcpyToHost,
         QuickGelu,
         ScatterElements,
     ]

diff --git a/onnx_array_api/reference/ops/op_memcpy_host.py b/onnx_array_api/reference/ops/op_memcpy_host.py
@@ -0,0 +1,11 @@
+from onnx.reference.op_run import OpRun
+
+
+class MemcpyFromHost(OpRun):
+    def _run(self, x):
+        return (x,)
+
+
+class MemcpyToHost(OpRun):
+    def _run(self, x):
+        return (x,)