if code coverage

Author: xadupre

_unittests/ut_reference/test_ort_evaluator.py

@@ -4,6 +4,7 @@
 import ml_dtypes
 from onnx import ModelProto, TensorProto
 from onnx.checker import check_model
+import onnx
 import onnx.helper as oh
 import onnx.numpy_helper as onh
 import torch
@@ -248,6 +249,71 @@ def test_init_torch_bfloat16(self):
         self.assertIsInstance(got[0], (torch.Tensor, np.ndarray))
         self.assertEqualArray(expected[0], got[0])
 
+    @hide_stdout()
+    def test_if(self):
+
+        def _mkv_(name):
+            value_info_proto = onnx.ValueInfoProto()
+            value_info_proto.name = name
+            return value_info_proto
+
+        model = oh.make_model(
+            oh.make_graph(
+                [
+                    oh.make_node("ReduceSum", ["X"], ["Xred"]),
+                    oh.make_node("Add", ["X", "two"], ["X0"]),
+                    oh.make_node("Add", ["X0", "zero"], ["X00"]),
+                    oh.make_node("CastLike", ["one", "Xred"], ["one_c"]),
+                    oh.make_node("Greater", ["Xred", "one_c"], ["cond"]),
+                    oh.make_node(
+                        "If",
+                        ["cond"],
+                        ["Z_c"],
+                        then_branch=oh.make_graph(
+                            [
+                                oh.make_node("Constant", [], ["two"], value_floats=[2.1]),
+                                oh.make_node("Add", ["X00", "two"], ["Y"]),
+                            ],
+                            "then",
+                            [],
+                            [_mkv_("Y")],
+                        ),
+                        else_branch=oh.make_graph(
+                            [
+                                oh.make_node("Constant", [], ["two"], value_floats=[2.2]),
+                                oh.make_node("Sub", ["X0", "two"], ["Y"]),
+                            ],
+                            "else",
+                            [],
+                            [_mkv_("Y")],
+                        ),
+                    ),
+                    oh.make_node("CastLike", ["Z_c", "X"], ["Z"]),
+                ],
+                "test",
+                [
+                    oh.make_tensor_value_info("X", TensorProto.FLOAT, ["N"]),
+                    oh.make_tensor_value_info("one", TensorProto.FLOAT, ["N"]),
+                ],
+                [oh.make_tensor_value_info("Z", TensorProto.UNDEFINED, ["N"])],
+                [
+                    onh.from_array(np.array([0], dtype=np.float32), name="zero"),
+                    onh.from_array(np.array([2], dtype=np.float32), name="two"),
+                ],
+            ),
+            opset_imports=[oh.make_operatorsetid("", 18)],
+            ir_version=10,
+        )
+        feeds = {
+            "X": np.array([1, 2, 3], dtype=np.float32),
+            "one": np.array([1], dtype=np.float32),
+        }
+        ref = ExtendedReferenceEvaluator(model, verbose=10)
+        expected = ref.run(None, feeds)[0]
+        sess = OnnxruntimeEvaluator(model, verbose=10)
+        got = sess.run(None, feeds)[0]
+        self.assertEqualArray(expected[0], got[0])
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)