fix plugs

Author: xadupre

CHANGELOGS.rst

@@ -4,6 +4,8 @@ Change Logs
 0.8.4
 +++++
 
+* :pr:`336`: implements versioned onnx plugs
+
 0.8.3
 +++++
 

onnx_diagnostic/export/api.py

@@ -149,6 +149,7 @@ def find_method(self, name: Any):
 
     if exporter in ("dynamo", "onnx-dynamo"):
         import os
+        from ..helpers import flatten_object
         import onnxscript.rewriter.ort_fusions as ort_fusions
 
         assert (
@@ -180,7 +181,12 @@ def find_method(self, name: Any):
             import onnx_ir as ir
             import onnx_ir.passes.common as common_passes
 
-            irfunctions = [ir.from_proto(plug.function_proto) for plug in onnx_plugs]
+            irfunctions = [
+                ir.from_proto(
+                    plug.get_function_proto(*flatten_object((args, kwargs), drop_keys=True))
+                )
+                for plug in onnx_plugs
+            ]
             for func in irfunctions:
                 epo.model.functions[func.identifier()] = func
             if inline:

onnx_diagnostic/export/onnx_plug.py

@@ -210,10 +210,10 @@ def _check_protos(self, params):
         # multiple protos
         assert all(
             self.n_inputs == len(v.input) for v in self._function_proto_versioned.values()
-        ), f"Output mismatch n_inputs={self.n_inputs} but one verion is wrong"
+        ), f"Output mismatch n_inputs={self.n_inputs} but one version is wrong"
         assert all(
             self.n_outputs == len(v.output) for v in self._function_proto_versioned.values()
-        ), f"Output mismatch n_outputs={self.n_outputs} but one verion is wrong"
+        ), f"Output mismatch n_outputs={self.n_outputs} but one version is wrong"
         assert all(
             v.domain == self.domain for v in self._function_proto_versioned.values()
         ), f"Function domain must be {self.domain!r} but it is different in one version"