Fix examples with scan

Author: xadupre

_unittests/ut_reference/test_onnxruntime_evaluator.py

@@ -0,0 +1,64 @@
+import unittest
+import onnx
+import torch
+import onnxruntime
+from onnx_diagnostic.ext_test_case import ExtTestCase
+from onnx_diagnostic.reference import OnnxruntimeEvaluator, ExtendedReferenceEvaluator
+
+try:
+    from experimental_experiment.torch_interpreter import to_onnx, ExportOptions
+except ImportError:
+    to_onnx = None
+
+
+class TestOnnxruntimeEvaluator(ExtTestCase):
+    def test_ort_eval_scan_cdist_add(self):
+
+        def dist(unused: torch.Tensor, x: torch.Tensor, samex: torch.Tensor):
+            sub = samex - x.reshape((1, -1))
+            sq = sub * sub
+            rd = torch.sqrt(sq.sum(axis=1))
+            # clone --> UnsupportedAliasMutationException:
+            # Combine_fn might be aliasing the input!
+            return [unused.clone(), rd]
+
+        class ScanModel(torch.nn.Module):
+            def forward(self, x):
+                z = torch.tensor([0], dtype=torch.float32)
+                y = x.clone()
+                out = torch.ops.higher_order.scan(dist, [z], [x], additional_inputs=[y])
+                return out[1]
+
+        x = torch.tensor([[1, 2, 3, -1], [4, 5, 6, -1], [7, 8, 9, -1]], dtype=torch.float32)
+        model = ScanModel()
+        expected = model(x)
+        onx = to_onnx(
+            model,
+            (x,),
+            optimize=True,
+            export_options=ExportOptions(decomposition_table="default", strict=False),
+            inline=False,
+        )
+        filename = self.get_dump_file("test_ort_eval_scan_cdist_add.onnx")
+        onnx.save(onx, filename)
+        inits = [i.name for i in onx.graph.initializer]
+        self.assertEqual(inits, ["c_lifted_tensor_0"])
+        name = onx.graph.input[0].name
+
+        sess = onnxruntime.InferenceSession(
+            onx.SerializeToString(), providers=["CPUExecutionProvider"]
+        )
+        got = sess.run(None, {name: x.numpy()})[0]
+        self.assertEqualArray(expected, got)
+
+        ref = ExtendedReferenceEvaluator(onx)
+        got = ref.run(None, {name: x.numpy()})[0]
+        self.assertEqualArray(expected, got)
+
+        orte = OnnxruntimeEvaluator(onx)
+        got = orte.run(None, {name: x.numpy()})[0]
+        self.assertEqualArray(expected, got)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

onnx_diagnostic/reference/evaluator.py

@@ -33,6 +33,7 @@
 from .ops.op_quick_gelu import QuickGelu
 from .ops.op_replace_zero import ReplaceZero
 from .ops.op_rotary import Rotary
+from .ops.op_scan import Scan
 from .ops.op_scatter_elements import ScatterElements
 from .ops.op_scatternd_of_shape import MaskedScatterNDOfShape, ScatterNDOfShape
 from .ops.op_simplified_layer_normalization import SimplifiedLayerNormalization
@@ -99,6 +100,7 @@ class ExtendedReferenceEvaluator(ReferenceEvaluator):
         QuickGelu,
         ReplaceZero,
         Rotary,
+        Scan,
         ScatterElements,
         ScatterNDOfShape,
         SimplifiedLayerNormalization,

onnx_diagnostic/reference/ops/op_scan.py

@@ -0,0 +1,65 @@
+import numpy as np
+from onnx.reference.ops.op_scan import Scan as _Scan
+
+
+class Scan(_Scan):
+
+    def need_context(self) -> bool:
+        """Tells the runtime if this node needs the context
+        (all the results produced so far) as it may silently access
+        one of them (operator Loop).
+        The default answer is `False`.
+        """
+        return True
+
+    def _run(
+        self,
+        *args,
+        context=None,
+        body=None,
+        num_scan_inputs=None,
+        scan_input_axes=None,
+        scan_input_directions=None,
+        scan_output_axes=None,
+        scan_output_directions=None,
+        attributes=None,
+    ):
+        (
+            num_loop_state_vars,
+            num_scan_outputs,
+            output_directions,
+            max_dir_out,
+            output_axes,
+            max_axe_out,
+            state_names_in,
+            state_names_out,
+            scan_names_in,
+            scan_names_out,
+            scan_values,
+            states,
+        ) = self._common_run_shape(*args)
+
+        max_iter = args[num_loop_state_vars].shape[self.input_axes_[0]]
+        results = [[] for _ in scan_names_out]  # type: ignore
+
+        for it in range(max_iter):
+            inputs = context.copy()
+            inputs.update(dict(zip(state_names_in, states)))
+            inputs.update({name: value[it] for name, value in zip(scan_names_in, scan_values)})
+
+            try:
+                outputs_list = self._run_body(inputs)  # type: ignore
+            except TypeError as e:
+                raise TypeError(
+                    f"Unable to call 'run' for type '{type(self.body)}'."  # type: ignore
+                ) from e
+
+            outputs = dict(zip(self.output_names, outputs_list))
+            states = [outputs[name] for name in state_names_out]
+            for i, name in enumerate(scan_names_out):
+                results[i].append(np.expand_dims(outputs[name], axis=0))
+
+        for res in results:
+            conc = np.vstack(res)
+            states.append(conc)
+        return self._check_and_fix_outputs(tuple(states))