add sumarry

sdpython · sdpython · Feb 6, 2024 · Feb 2, 2024 · Feb 2, 2024 · Feb 2, 2024
commit fe32241f9d4b079a701a52ddc19d8621b0175f57
diff --git a/_unittests/ut_reference/test_evaluator_yield.py b/_unittests/ut_reference/test_evaluator_yield.py
@@ -62,31 +62,37 @@ def test_evaluator_yield(self):
                 ResultType.INPUT,
                 "A",
                 np.array([[0.0, 1.0], [2.0, 3.0]], dtype=np.float32),
+                None,
             ),
             (
                 ResultType.INPUT,
                 "B",
                 np.array([[0.0, 1.0], [2.0, 3.0]], dtype=np.float32),
+                None,
             ),
             (
                 ResultType.INPUT,
                 "X",
                 np.array([[0.0, 1.0], [2.0, 3.0]], dtype=np.float32),
+                None,
             ),
             (
                 ResultType.RESULT,
                 "Y1",
                 np.array([[2.0, 4.0], [8.0, 14.0]], dtype=np.float32),
+                "LinearRegression",
             ),
             (
                 ResultType.RESULT,
                 "Y",
                 np.array([[2.0, 4.0], [8.0, 14.0]], dtype=np.float32),
+                "Abs",
             ),
             (
                 ResultType.OUTPUT,
                 "Y",
                 np.array([[2.0, 4.0], [8.0, 14.0]], dtype=np.float32),
+                None,
             ),
         ]
         self.assertEqual(len(expected), len(results))
@@ -95,6 +101,73 @@ def test_evaluator_yield(self):
             self.assertEqual(a[0], b[0])
             self.assertEqual(a[1], b[1])
             self.assertEqual(a[2].tolist(), b[2].tolist())
+            self.assertEqual(a[3], b[3])
+
+    def test_evaluator_yield_summary(self):
+        new_domain = "custom_domain"
+        opset_imports = [make_opsetid("", 14), make_opsetid(new_domain, 1)]
+
+        node1 = make_node("MatMul", ["X", "A"], ["XA"])
+        node2 = make_node("Add", ["XA", "B"], ["Y"])
+
+        linear_regression = make_function(
+            new_domain,
+            "LinearRegression",
+            ["X", "A", "B"],
+            ["Y"],
+            [node1, node2],
+            opset_imports,
+            [],
+        )
+
+        X = make_tensor_value_info("X", TensorProto.FLOAT, [None, None])
+        A = make_tensor_value_info("A", TensorProto.FLOAT, [None, None])
+        B = make_tensor_value_info("B", TensorProto.FLOAT, [None, None])
+        Y = make_tensor_value_info("Y", TensorProto.FLOAT, None)
+
+        graph = make_graph(
+            [
+                make_node(
+                    "LinearRegression", ["X", "A", "B"], ["Y1"], domain=new_domain
+                ),
+                make_node("Abs", ["Y1"], ["Y"]),
+            ],
+            "example",
+            [X, A, B],
+            [Y],
+        )
+
+        onnx_model = make_model(
+            graph, opset_imports=opset_imports, functions=[linear_regression]
+        )
+
+        cst = np.arange(4).reshape((-1, 2)).astype(np.float32)
+        yield_eval = YieldEvaluator(onnx_model)
+        results = list(
+            yield_eval.enumerate_summarized(None, {"A": cst, "B": cst, "X": cst})
+        )
+        expected = [
+            (ResultType.INPUT, np.dtype("float32"), (2, 2), "ABCD", None),
+            (ResultType.INPUT, np.dtype("float32"), (2, 2), "ABCD", None),
+            (ResultType.INPUT, np.dtype("float32"), (2, 2), "ABCD", None),
+            (
+                ResultType.RESULT,
+                np.dtype("float32"),
+                (2, 2),
+                "CEIO",
+                "LinearRegression",
+            ),
+            (ResultType.RESULT, np.dtype("float32"), (2, 2), "CEIO", "Abs"),
+            (ResultType.OUTPUT, np.dtype("float32"), (2, 2), "CEIO", None),
+        ]
+        self.assertEqual(len(expected), len(results))
+        for a, b in zip(expected, results):
+            self.assertEqual(len(a), len(b))
+            self.assertEqual(a[0], b[0])
+            self.assertEqual(a[1], b[1])
+            self.assertEqual(a[2], b[2])
+            self.assertEqual(a[3], b[3])
+            self.assertEqual(a[4], b[4])
 
 
 if __name__ == "__main__":

diff --git a/onnx_array_api/reference/evaluator_yield.py b/onnx_array_api/reference/evaluator_yield.py
@@ -1,5 +1,6 @@
 from typing import Any, Dict, List, Iterator, Optional, Tuple
 from enum import IntEnum
+import numpy as np
 from onnx import ModelProto
 from .evaluator import ExtendedReferenceEvaluator
 
@@ -41,14 +42,14 @@ def enumerate_results(
         feed_inputs: Optional[Dict[str, Any]] = None,
     ) -> Iterator[Tuple[ResultType, str, Any]]:
         """
-        Executes the onnx model.
+        Executes the onnx model and enumerate all the intermediate results.
 
         Args:
             output_names: requested outputs by names, None for all
             feed_inputs: dictionary `{ input name: input value }`
 
         Returns:
-            iterator on tuple(result kind, name, value)
+            iterator on tuple(result kind, name, value, node.op_type or None)
         """
         assert isinstance(self.evaluator, ExtendedReferenceEvaluator), (
             f"This implementation only works with "
@@ -63,10 +64,10 @@ def enumerate_results(
         results.update(feed_inputs)
         # step 0: initializer
         for k, v in self.evaluator.rt_inits_.items():
-            yield ResultType.INITIALIZER, k, v
+            yield ResultType.INITIALIZER, k, v, None
         # step 1: inputs
         for k, v in feed_inputs.items():
-            yield ResultType.INPUT, k, v
+            yield ResultType.INPUT, k, v, None
 
         # step 2: execute nodes
         for node in self.evaluator.rt_nodes_:
@@ -86,7 +87,7 @@ def enumerate_results(
             else:
                 outputs = node.run(*inputs, **linked_attributes)
             for name, value in zip(node.output, outputs):
-                yield ResultType.RESULT, name, value
+                yield ResultType.RESULT, name, value, node.op_type
                 results[name] = value
 
         # step 3: outputs
@@ -95,4 +96,50 @@ def enumerate_results(
                 raise RuntimeError(
                     f"Unable to find output name {name!r} in {sorted(results)}, proto is\n{self.proto_}"
                 )
-            yield ResultType.OUTPUT, name, results[name]
+            yield ResultType.OUTPUT, name, results[name], None
+
+    def enumerate_summarized(
+        self,
+        output_names: Optional[List[str]] = None,
+        feed_inputs: Optional[Dict[str, Any]] = None,
+    ) -> Iterator[Tuple[ResultType, str, Any]]:
+        """
+        Executes the onnx model and enumerate intermediate results without their names.
+
+        Args:
+            output_names: requested outputs by names, None for all
+            feed_inputs: dictionary `{ input name: input value }`
+
+        Returns:
+            iterator on tuple(result kind, node.type, dtype, shape, value)
+        """
+        for kind, name, value, op_type in self.enumerate_results(
+            output_names, feed_inputs
+        ):
+            summary = self.make_summary(value)
+            yield kind, value.dtype, value.shape, summary, op_type
+
+    def make_summary(self, value: Any, length: int = 4, modulo: int = 26) -> str:
+        """
+        Create a short string summarizing the value (discretization).
+
+        :param value: array
+        :param length: number of value to produce
+        :param module: discretization parameter
+        :return: short string
+        """
+        value4 = np.zeros(4, dtype=np.float64)
+        if value.size <= length:
+            value4[: value.size] = value.flatten().astype(np.float64)
+        else:
+            if value.size % length != 2:
+                value2 = np.zeros(
+                    value.size + length - value.size % length, dtype=np.float64
+                )
+                value2[: value.size] = value.flatten().astype(np.float64)
+            else:
+                value2 = value.flatten().astype(np.float64)
+            value4 = value2.reshape((4, -1)).mean(axis=1)
+        value4i = value4.astype(np.int64) % modulo
+        s = "".join([chr(65 + i) for i in value4i])
+        return s