fix test

Author: xadupre

_unittests/ut_export/test_control_flow.py

@@ -66,7 +66,10 @@ def body(i, x):
         ep = torch.export.export(
             model, (n_iter, x), dynamic_shapes=({}, ({0: torch.export.Dim.DYNAMIC}))
         )
-        self.assertIn("torch.ops.onnx_higher_ops.loop_for_body_", str(ep))
+        self.assertIn(
+            "torch.ops.onnx_higher_ops.loop_for_TestControlFlow_test_loop_one_custom_L_Model_forward_L_body_",
+            str(ep),
+        )
 
         onx = to_onnx(
             model,
@@ -97,7 +100,10 @@ def body(i, x):
         ep = torch.export.export(
             model, (n_iter, x), dynamic_shapes=({}, ({0: torch.export.Dim.DYNAMIC}))
         )
-        self.assertIn("torch.ops.onnx_higher_ops.loop_for_body_", str(ep))
+        self.assertIn(
+            "torch.ops.onnx_higher_ops.loop_for_TestControlFlow_test_loop_one_custom_different_opset_L_Model_forward_L_body_",
+            str(ep),
+        )
 
         onx = to_onnx(
             model,
@@ -132,7 +138,10 @@ def body(i, x):
         ep = torch.export.export(
             model, (n_iter, x), dynamic_shapes=({}, ({0: torch.export.Dim.DYNAMIC}))
         )
-        self.assertIn("torch.ops.onnx_higher_ops.loop_for_body_", str(ep))
+        self.assertIn(
+            "torch.ops.onnx_higher_ops.loop_for_TestControlFlow_test_loop_two_custom_L_Model_forward_L_body_",
+            str(ep),
+        )
 
         onx = to_onnx(
             model,
@@ -164,7 +173,10 @@ def body(i, x):
         ep = torch.export.export(
             model, (n_iter, x), dynamic_shapes=({}, ({0: torch.export.Dim.DYNAMIC}))
         )
-        self.assertIn("torch.ops.onnx_higher_ops.loop_for_body_", str(ep))
+        self.assertIn(
+            "torch.ops.onnx_higher_ops.loop_for_TestControlFlow_test_loop_two_custom_reduction_dim_L_Model_forward_L_body_",
+            str(ep),
+        )
 
         onx = to_onnx(
             model,

_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
+from onnx_diagnostic.ext_test_case import ExtTestCase, has_torch
 from onnx_diagnostic.export.onnx_plug import EagerDirectReplacementWithOnnx
 from onnx_diagnostic.export.api import to_onnx
 
@@ -85,6 +85,8 @@ def forward(self, x):
             )
             self.assert_onnx_disc("test_onnx_plug_export_custom", onx.model_proto, model, (x,))
 
+        if not has_torch("2.9"):
+            raise unittest.SkipTest("onnx-dynamo + custom op not fully working on 2.8")
         with self.subTest(exporter="onnx-dynamo"):
             onx = to_onnx(
                 model,

onnx_diagnostic/export/api.py

@@ -20,6 +20,7 @@ def to_onnx(
     optimize: bool = True,
     use_control_flow_dispatcher: bool = False,
     onnx_plugs: Optional[List[EagerDirectReplacementWithOnnx]] = None,
+    inline: bool = True,
 ) -> Any:
     """
     Common API for exporters. By default, the models are optimized to use the
@@ -44,6 +45,7 @@ def to_onnx(
     :param use_control_flow_dispatcher: use the dispatcher created to supported
         custom loops (see :func:`onnx_diagnostic.export.control_flow.loop_for`)
     :param onnx_plugs: the code was modified to replace some parts with onnx translation
+    :param inline: inline local functions
     :return: the output of the selected exporter, usually a structure including
         an onnx model
 
@@ -63,6 +65,11 @@ def to_onnx(
     :func:`onnx_diagnostic.export.control_flow.loop_for` or
     :class:`onnx_diagnostic.export.onnx_plug.EagerDirectReplacementWithOnnx`.
     """
+    if exporter_kwargs and "inline" in exporter_kwargs:
+        assert (
+            inline == exporter_kwargs["inline"]
+        ), f"Mismatch between inline={inline} and exporter_kwargs={exporter_kwargs}"
+        exporter_kwargs.pop("inline")
     if exporter == "custom":
         from experimental_experiment.torch_interpreter import (
             to_onnx as _to_onnx,
@@ -90,7 +97,7 @@ def __init__(self):
                     super().__init__({})
 
             main_dispatcher = MainDispatcher()
-            if control_flow_dispatcher:
+            if use_control_flow_dispatcher:
                 main_dispatcher.registered_functions.update(
                     control_flow_dispatcher.registered_functions
                 )
@@ -99,7 +106,6 @@ def __init__(self):
                     main_dispatcher.registered_functions[plug.target_name] = (
                         plug.custom_converter()
                     )
-
         else:
             main_dispatcher = None
 
@@ -117,8 +123,9 @@ def __init__(self):
             output_dynamic_shapes=output_dynamic_shapes,
             export_options=ExportOptions(save_ep=save_ep),
             options=options,
-            **(exporter_kwargs or {}),
+            inline=inline,
             dispatcher=main_dispatcher,
+            **(exporter_kwargs or {}),
         )
 
     if exporter in ("dynamo", "onnx-dynamo"):
@@ -147,7 +154,21 @@ def __init__(self):
             custom_translation_table=custom_translation_table,
             **(exporter_kwargs or {}),
         )
-        if optimize:
+        if not inline and optimize:
+            ort_fusions.optimize_for_ort(epo.model)
+
+        if onnx_plugs:
+            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]
+            for func in irfunctions:
+                epo.model.functions[func.identifier()] = func
+            if inline:
+                common_passes.InlinePass()(epo.model)
+                common_passes.RemoveUnusedOpsetsPass()(epo.model)
+
+        if inline and optimize:
             ort_fusions.optimize_for_ort(epo.model)
         if filename:
             epo.save(filename, external_data=True)

onnx_diagnostic/export/control_flow.py

@@ -134,7 +134,8 @@ def make_custom_loop_for(
     assert body_outputs is not None, "body_outputs cannot be None"
     srank = "_".join("x".join(map(str, s.shape)) for s in body_outputs)
     sred = "x".join(map(str, reduction_dim)) if reduction_dim else ""
-    name = f"loop_for_{body_fn.__name__}_{id(body_fn)}_{srank}_{sred}"
+    full_name = body_fn.__qualname__.replace("<locals>", "L").replace(".", "_")
+    name = f"loop_for_{full_name}_{srank}_{sred}"
     if name in _REGISTERED_SCHEMA:
         return name, _REGISTERED_SCHEMA[name][0]
     sig = inspect.signature(body_fn)

onnx_diagnostic/export/onnx_plug.py

@@ -49,6 +49,88 @@ class EagerDirectReplacementWithOnnx:
     :param n_outputs: same for the number of outputs,
         only tensors must be counted
     :param name: the name of the custom op, the function name if not specified
+
+    Here is an example:
+
+    .. runpython::
+        :showcode:
+
+        import onnx.helper as oh
+        import torch
+        from onnx_diagnostic.helpers.onnx_helper import pretty_onnx
+        from onnx_diagnostic.export.onnx_plug import EagerDirectReplacementWithOnnx
+        from onnx_diagnostic.export.api import to_onnx
+        from onnx_diagnostic.torch_export_patches.patch_inputs import use_dyn_not_str
+
+
+        def demo_customsub(x, y):
+            return x - y
+
+
+        def demo_customsub_shape(x, y):
+            return torch.empty(torch.broadcast_shapes(x.shape, y.shape), dtype=x.dtype)
+
+
+        def make_function_proto():
+            return oh.make_function(
+                "onnx_plug",
+                "demo_customsub",
+                ["x", "y"],
+                ["z"],
+                [oh.make_node("Sub", ["x", "y"], ["z"])],
+                opset_imports=[oh.make_opsetid("", 22)],
+            )
+
+
+        class Model(torch.nn.Module):
+            def forward(self, x):
+                y = x.sum(axis=1, keepdim=True)
+                d = torch.ops.onnx_plug.demo_customsub(x, y)
+                return torch.abs(d)
+
+
+        replacements = [
+            EagerDirectReplacementWithOnnx(
+                demo_customsub, demo_customsub_shape, make_function_proto(), 2, 1
+            )
+        ]
+
+        x = torch.randn((3, 4), dtype=torch.float32)
+        model = Model()
+        ds = ({0: "d1", 1: "d2"},)
+
+        # The exported program shows a custom op.
+        ep = torch.export.export(model, (x,), dynamic_shapes=use_dyn_not_str(ds))
+        print("ep")
+
+        # As the exporter knows how the replace this custom op.
+        # Let's export.
+
+        onx = to_onnx(
+            model,
+            (x,),
+            dynamic_shapes=ds,
+            exporter="custom",
+            onnx_plugs=replacements,
+            target_opset=22,
+            inline=False,
+        ).model_proto
+
+        print(pretty_onnx(onx))
+
+        # And with :func:`torch.onnx.export`:
+
+        onx = to_onnx(
+            model,
+            (x,),
+            dynamic_shapes=ds,
+            exporter="onnx-dynamo",
+            onnx_plugs=replacements,
+            target_opset=22,
+            inline=False,
+        ).model_proto
+
+        print(pretty_onnx(onx))
     """
 
     def __init__(