if

Author: xadupre

_unittests/ut_reference/test_torch_onnx_evaluator.py

@@ -897,6 +897,65 @@ def test_local_function(self):
             atol=1e-6,
         )
 
+    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", [], ["t2"], value_floats=[2.1]),
+                                oh.make_node("Add", ["X00", "t2"], ["Y"]),
+                            ],
+                            "then",
+                            [],
+                            [_mkv_("Y")],
+                        ),
+                        else_branch=oh.make_graph(
+                            [
+                                oh.make_node("Constant", [], ["t2"], value_floats=[2.2]),
+                                oh.make_node("Sub", ["X0", "t2"], ["Y"]),
+                            ],
+                            "else",
+                            [],
+                            [_mkv_("Y")],
+                        ),
+                    ),
+                    oh.make_node("CastLike", ["Z_c", "X"], ["Z"]),
+                ],
+                "test",
+                [
+                    oh.make_tensor_value_info("X", TFLOAT, ["N"]),
+                    oh.make_tensor_value_info("one", TFLOAT, ["N"]),
+                ],
+                [oh.make_tensor_value_info("Z", onnx.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,
+        )
+        self._finalize_test(
+            model,
+            torch.tensor([1, 2, 3], dtype=torch.float32),
+            torch.tensor([1], dtype=torch.float32),
+        )
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/reference/torch_ops/__init__.py

@@ -17,7 +17,7 @@
 )
 from .generator_ops import Range_11
 from .nn_ops import LayerNormalization_17, Softmax_13, Tanh_6
-from .other_ops import Cast_6, Concat_1, Transpose_1, Trilu_14, Where_9
+from .other_ops import Cast_6, CastLike_15, Concat_1, Transpose_1, Trilu_14, Where_9
 from .reduce_ops import ReduceMax_18, ReduceMean_18, ReduceMin_18, ReduceSum_18
 from .shape_ops import (
     ConstantOfShape_9,

onnx_diagnostic/reference/torch_ops/other_ops.py

@@ -13,11 +13,25 @@ def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
         to = self.get_attribute_int(node, "to", 0)
         assert isinstance(to, int), f"Unexpected value for attribute to={to!r}"
         self.to = onnx_dtype_to_torch_dtype(to)
+        self.saturate = self.get_attribute_int(node, "saturate", 1)
+        assert self.saturate == 1, f"saturate={self.saturate} not implemented for Cast"
 
     def run(self, data: OpRunValue) -> OpRunValue:
         return OpRunValue(data.tensor.to(self.to))
 
 
+class CastLike_15(OpRun):
+    "Cast"
+
+    def __init__(self, node: onnx.NodeProto, version: Optional[int] = None):
+        super().__init__(node, version)
+        self.saturate = self.get_attribute_int(node, "saturate", 1)
+        assert self.saturate == 1, f"saturate={self.saturate} not implemented for CastLike"
+
+    def run(self, data: OpRunValue, like: OpRunValue) -> OpRunValue:
+        return OpRunValue(data.tensor.to(like.tensor.dtype))
+
+
 class Concat_1(OpRun):
     "Concat"