Export evaluator type in compare_onnx_execution

xadupre · xadupre · commit b11db3c1475b · 2025-01-06T16:28:14.000+01:00
diff --git a/LICENSE.txt b/LICENSE.txt
@@ -1,4 +1,4 @@
-﻿Copyright (c) 2023-2024, Xavier Dupré
+﻿Copyright (c) 2023-2025, Xavier Dupré
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal
diff --git a/onnx_array_api/reference/evaluator_yield.py b/onnx_array_api/reference/evaluator_yield.py
@@ -3,6 +3,7 @@
 from enum import IntEnum
 import numpy as np
 from onnx import ModelProto, TensorProto, ValueInfoProto, load
+from onnx.reference import ReferenceEvaluator
 from onnx.helper import tensor_dtype_to_np_dtype
 from onnx.shape_inference import infer_shapes
 from . import to_array_extended
@@ -166,9 +167,9 @@ def enumerate_results(
         Returns:
             iterator on tuple(result kind, name, value, node.op_type or None)
         """
-        assert isinstance(self.evaluator, ExtendedReferenceEvaluator), (
+        assert isinstance(self.evaluator, ReferenceEvaluator), (
             f"This implementation only works with "
-            f"ExtendedReferenceEvaluator not {type(self.evaluator)}"
+            f"ReferenceEvaluator not {type(self.evaluator)}"
         )
         attributes = {}
         if output_names is None:
@@ -595,6 +596,7 @@ def compare_onnx_execution(
     raise_exc: bool = True,
     mode: str = "execute",
     keep_tensor: bool = False,
+    cls: Optional[type[ReferenceEvaluator]] = None,
 ) -> Tuple[List[ResultExecution], List[ResultExecution], List[Tuple[int, int]]]:
     """
     Compares the execution of two onnx models.
@@ -611,6 +613,7 @@ def compare_onnx_execution(
     :param mode: the model should be executed but the function can be executed
         but the comparison may append on nodes only
     :param keep_tensor: keeps the tensor in order to compute a precise distance
+    :param cls: evaluator class to use
     :return: four results, a sequence of results
         for the first model and the second model,
         the alignment between the two, DistanceExecution
@@ -634,15 +637,15 @@ def compare_onnx_execution(
             print(f"[compare_onnx_execution] execute with {len(inputs)} inputs")
             print("[compare_onnx_execution] execute first model")
         res1 = list(
-            YieldEvaluator(model1).enumerate_summarized(
+            YieldEvaluator(model1, cls=cls).enumerate_summarized(
                 None, feeds1, raise_exc=raise_exc, keep_tensor=keep_tensor
             )
         )
         if verbose:
             print(f"[compare_onnx_execution] got {len(res1)} results")
             print("[compare_onnx_execution] execute second model")
         res2 = list(
-            YieldEvaluator(model2).enumerate_summarized(
+            YieldEvaluator(model2, cls=cls).enumerate_summarized(
                 None, feeds2, raise_exc=raise_exc, keep_tensor=keep_tensor
             )
         )

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-Copyright (c) 2023-2024, Xavier Dupré`
	`1`	`+Copyright (c) 2023-2025, Xavier Dupré`
`2`	`2`
`3`	`3`	`Permission is hereby granted, free of charge, to any person obtaining a copy`
`4`	`4`	`of this software and associated documentation files (the "Software"), to deal`