fix bug in onnxruntime

Author: xadupre

_unittests/ut_reference/test_onnxruntime_evaluator.py

@@ -1,8 +1,11 @@
 import unittest
+import numpy as np
 import onnx
+import onnx.helper as oh
 import torch
 import onnxruntime
 from onnx_diagnostic.ext_test_case import ExtTestCase
+from onnx_diagnostic.helpers.onnx_helper import from_array_extended
 from onnx_diagnostic.reference import OnnxruntimeEvaluator, ExtendedReferenceEvaluator
 
 try:
@@ -96,6 +99,52 @@ def false_fn(x, y):
             for e, g in zip(expected, got):
                 self.assertEqualArray(e, g, atol=1e-5)
 
+    def test_constant_bool(self):
+        node = oh.make_node(
+            "Constant",
+            [],
+            ["cbool"],
+            value=from_array_extended(np.array(True, dtype=np.bool_)),
+        )
+        ref = ExtendedReferenceEvaluator(node)
+        got = ref.run(None, {})[0]
+        self.assertEqual(got.dtype, np.bool_)
+        self.assertEqual(got, True)
+        ref = OnnxruntimeEvaluator(node)
+        got = ref.run(None, {})[0]
+        self.assertEqual(len(ref._cache), 1)
+        values = list(ref._cache.values())
+        _, sess = values[0]
+        got2 = sess.run(None, {})[0]
+        self.assertIn(got2.dtype, (torch.bool, np.bool_))
+        self.assertEqual(got2, True)
+
+        self.assertIn(got.dtype, (torch.bool, np.bool_))
+        self.assertEqual(got, True)
+
+    def test_constant_bool_array(self):
+        node = oh.make_node(
+            "Constant",
+            [],
+            ["cbool"],
+            value=from_array_extended(np.array([True], dtype=np.bool_)),
+        )
+        ref = ExtendedReferenceEvaluator(node)
+        got = ref.run(None, {})[0]
+        self.assertEqual(got.dtype, np.bool_)
+        self.assertEqual(got[0], True)
+        ref = OnnxruntimeEvaluator(node)
+        got = ref.run(None, {})[0]
+        self.assertEqual(len(ref._cache), 1)
+        values = list(ref._cache.values())
+        _, sess = values[0]
+        got2 = sess.run(None, {})[0]
+        self.assertIn(got2.dtype, (torch.bool, np.bool_))
+        self.assertEqual(got2[0], True)
+
+        self.assertIn(got.dtype, (torch.bool, np.bool_))
+        self.assertEqual(got[0], True)
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/helpers/ort_session.py

@@ -154,6 +154,7 @@ def __init__(
                 )
 
         self._torch_from_dlpack = _from_dlpack
+        self.sess_bool_outputs = [i.type == "tensor(bool)" for i in sess.get_outputs()]
 
     def run(
         self,
@@ -166,7 +167,19 @@ def run(
         ort_outputs = self.sess._sess.run_with_ort_values(
             feeds, output_names or self.output_names, self.run_options
         )
-        return ort_outputs
+        return self._post_process_inplace(ort_outputs)
+
+    def _post_process_inplace(self, outputs):
+        for i in range(len(outputs)):
+            o = outputs[i]
+            if self.sess_bool_outputs[i]:
+                if isinstance(o, np.ndarray):
+                    if o.dtype != np.bool_:
+                        outputs[i] = o.astype(np.bool_)
+                else:
+                    if o.dtype != torch.bool:
+                        outputs[i] = o.to(torch.bool)
+        return outputs
 
 
 class InferenceSessionForNumpy(_InferenceSession):
@@ -221,7 +234,7 @@ def run(
         """Calls :meth:`onnxruntime.InferenceSession.run`."""
         # sess.run does not support blfoat16
         # res = self.sess.run(output_names, feeds)
-        return list(self.run_dlpack(output_names, feeds))
+        return self._post_process_inplace(list(self.run_dlpack(output_names, feeds)))
 
     def run_dlpack(
         self, output_names: Optional[List[str]], feeds: Dict[str, npt.ArrayLike]
@@ -421,7 +434,7 @@ def run(  # type: ignore
         if self.use_training_api:
             inputs = [feeds[i] for i in self.input_names]
             return self.run_training_api(*inputs, output_names=output_names)
-        return self.run_dlpack(output_names, feeds)
+        return self._post_process_inplace(list(self.run_dlpack(output_names, feeds)))
 
     def run_training_api(
         self, *inputs, output_names: Optional[List[str]] = None

onnx_diagnostic/reference/ort_evaluator.py

@@ -10,6 +10,8 @@
     ValueInfoProto,
     helper as oh,
     load,
+    save as onnx_save,
+    shape_inference as shi,
 )
 from onnx.defs import onnx_opset_version
 import onnxruntime
@@ -20,6 +22,7 @@
     InferenceSessionForNumpy,
     _InferenceSession,
 )
+from .evaluator import ExtendedReferenceEvaluator
 
 PROTO = (FunctionProto, ModelProto, GraphProto, NodeProto)
 Proto = Union[FunctionProto, ModelProto, GraphProto, NodeProto]
@@ -304,7 +307,7 @@ def _make_model_proto(
         return onx
 
     @classmethod
-    def _get_hidden_inputs(self, graph: GraphProto) -> Set[str, Any]:
+    def _get_hidden_inputs(self, graph: GraphProto) -> Set[str]:
         """
         Returns the hidden inputs (inputs coming from an upper context)
         used by a subgraph.
@@ -331,18 +334,34 @@ def _get_sess(
             onx = node
         else:
             assert isinstance(node, NodeProto), f"Unexpected type {type(node)} for node"
-            unique_names = set()
-            vinputs = []
-            for i, it in zip(node.input, inputs):
-                if i == "" or i in unique_names:
-                    continue
-                unique_names.add(i)
-                value = oh.make_tensor_value_info(i, dtype_to_tensor_dtype(it.dtype), it.shape)
-                vinputs.append(value)
+            if node.op_type == "Constant":
+                # We force the type to be a boolean.
+                ref = ExtendedReferenceEvaluator(node)
+                cst = ref.run(None, {})[0]
+                vinputs = []
+                voutputs = [
+                    oh.make_tensor_value_info(
+                        node.output[0], dtype_to_tensor_dtype(cst.dtype), cst.shape
+                    )
+                ]
+            else:
+                unique_names = set()
+                vinputs = []
+                for i, it in zip(node.input, inputs):
+                    if i == "" or i in unique_names:
+                        continue
+                    unique_names.add(i)
+                    value = oh.make_tensor_value_info(
+                        i, dtype_to_tensor_dtype(it.dtype), it.shape
+                    )
+                    vinputs.append(value)
+
+                # no need to run shape inference
+                voutputs = [oh.make_value_info(o, TypeProto()) for o in node.output]
 
-            # no need to run shape inference
-            voutputs = [oh.make_value_info(o, TypeProto()) for o in node.output]
             onx = self._make_model_proto([node], vinputs, voutputs)
+            # That helps fixing bugs.
+            onx = shi.infer_shapes(onx)
 
         cls = (
             InferenceSessionForNumpy
@@ -356,6 +375,7 @@ def _get_sess(
             onnxruntime.capi.onnxruntime_pybind11_state.InvalidGraph,
             onnxruntime.capi.onnxruntime_pybind11_state.InvalidArgument,
         ) as e:
+            onnx_save(onx, "_debug_OnnxruntimeEvaluator_last_failure.onnx")
             raise RuntimeError(
                 f"Unable to infer a session with inputs\n{string_type(inputs)}"
                 f"\ndue to {e}\n{pretty_onnx(onx)}"
@@ -461,8 +481,7 @@ def _run_if(
             self._cache[key] = onx, sess = self._get_sess_if(node, name, inputs, results)
 
         assert hasattr(sess, "run"), f"Missing method run for type {type(sess)}"
-        input_names = [i.name for i in sess.get_inputs()]
-        feeds = {name: results[name] for name in input_names}
+        feeds = {name: results[name] for name in sess.input_names}
         outputs = sess.run(None, feeds)
         assert isinstance(outputs, list), f"Unexpected type for outputs {type(outputs)}"
         return outputs
@@ -515,8 +534,7 @@ def _run_scan(
             self._cache[key] = onx, sess = self._get_sess_scan(node, name, inputs, results)
 
         assert hasattr(sess, "run"), f"Missing method run for type {type(sess)}"
-        input_names = [i.name for i in sess.get_inputs()]
-        feeds = {name: results[name] for name in input_names}
+        feeds = {name: results[name] for name in sess.input_names}
         outputs = sess.run(None, feeds)
         assert isinstance(outputs, list), f"Unexpected type for outputs {type(outputs)}"
         return outputs