fix

Author: xadupre

_unittests/ut_export/test_onnx_plug.py

@@ -1,7 +1,7 @@
 import unittest
 import onnx.helper as oh
 import torch
-from onnx_diagnostic.ext_test_case import ExtTestCase, has_torch
+from onnx_diagnostic.ext_test_case import ExtTestCase, has_torch, hide_stdout, ignore_warnings
 from onnx_diagnostic.export.onnx_plug import EagerDirectReplacementWithOnnx
 from onnx_diagnostic.export.api import to_onnx
 
@@ -36,6 +36,8 @@ def make_function_proto():
         self.assertEqual(len(res.diffs), 1)
         self.assertEqual(res.diffs[0]["abs"], 0)
 
+    @hide_stdout()
+    @ignore_warnings(FutureWarning)
     def test_onnx_plug_export(self):
         def _test_customsub(x, y):
             return x - y
@@ -61,7 +63,7 @@ def forward(self, x):
 
         replacements = [
             EagerDirectReplacementWithOnnx(
-                _test_customsub, _test_customsub_shape, make_function_proto(), 2, 1
+                _test_customsub, _test_customsub_shape, make_function_proto(), 2, 1, verbose=1
             )
         ]
 

onnx_diagnostic/export/onnx_plug.py

@@ -50,6 +50,7 @@ class EagerDirectReplacementWithOnnx:
         only tensors must be counted
     :param name: the name of the custom op, the function name if not specified
     :param kwargs: constants
+    :param verbose: verbose level
 
     Here is an example:
 
@@ -143,6 +144,7 @@ def __init__(
         n_outputs: Optional[int] = None,
         name: Optional[str] = None,
         kwargs: Optional[Dict[str, Union[int, float]]] = None,
+        verbose: int = 0,
     ):
         assert isinstance(
             function_proto, onnx.FunctionProto
@@ -154,9 +156,13 @@ def __init__(
         self.function_proto = function_proto
         self.n_inputs = n_inputs
         self.n_outputs = n_outputs
-        self.name = name or eager_fn.__qualname__.replace("<locals>", "L").replace(
-            "<lambda>", "l"
-        ).replace(".", "_")
+        self.name = name or (
+            eager_fn.__name__
+            if "<" not in eager_fn.__name__
+            else eager_fn.__qualname__.replace("<locals>", "L")
+            .replace("<lambda>", "l")
+            .replace(".", "_")
+        )
         self.kwargs = kwargs
         assert kwargs is None or all(isinstance(v, (int, float)) for v in kwargs.values()), (
             f"Only int or floats are allowed for kwargs={kwargs}, one of them "
@@ -179,7 +185,8 @@ def __init__(
             function_proto.domain == self.domain
         ), f"Function domain must be {self.domain!r} but it is {function_proto.domain!r}"
         self.arg_names = params
-        self.custom_op = self._registers()
+        self.verbose = verbose
+        self.custom_op = self._register()
 
     @property
     def domain(self) -> str:
@@ -202,12 +209,18 @@ def __call__(self, *args):
             return self.torch_op(*args)
         return self.eager_fn(*args)
 
-    def _registers(self):
+    def _register(self):
         """Registers the custom op."""
         inputs = ", ".join([f"Tensor {p}" for p in self.arg_names])
         schema = f"({inputs}) -> Tensor"
         if self.n_outputs > 1:
             schema += "[]"
+        if self.verbose:
+            print(
+                f"[EagerDirectReplacementWithOnnx._register] "
+                f"'torch.ops.{self.domain}.{self.name}"
+            )
+            print(f"[EagerDirectReplacementWithOnnx._register] schema={schema}")
         custom_def = torch.library.CustomOpDef(self.domain, self.name, schema, self.eager_fn)
         custom_def.register_kernel(None)(self.eager_fn)
         custom_def._abstract_fn = self.shape_fn