Check custom (#64)

* Check custom

* doc

Author: xadupre

_doc/conf.py

@@ -187,6 +187,7 @@
     "ExecuTorch": "https://pytorch.org/executorch/stable/intro-overview.html",
     "ExecuTorch Runtime Python API Reference": "https://pytorch.org/executorch/stable/runtime-python-api-reference.html",
     "ExecuTorch Tutorial": "https://pytorch.org/executorch/stable/tutorials/export-to-executorch-tutorial.html",
+    "experimental-experiment": "https://sdpython.github.io/doc/experimental-experiment/dev/",
     "JIT": "https://en.wikipedia.org/wiki/Just-in-time_compilation",
     "FunctionProto": "https://onnx.ai/onnx/api/classes.html#functionproto",
     "graph break": "https://pytorch.org/docs/stable/torch.compiler_faq.html#graph-breaks",

_unittests/ut_torch_models/test_test_helpers.py

@@ -7,6 +7,7 @@
     hide_stdout,
     ignore_warnings,
     requires_torch,
+    requires_experimental,
 )
 from onnx_diagnostic.torch_models.test_helper import (
     get_inputs_for_task,
@@ -62,16 +63,42 @@ def test_validate_model_export(self):
     @requires_torch("2.7")
     @hide_stdout()
     @ignore_warnings(FutureWarning)
-    def test_validate_model_onnx(self):
+    def test_validate_model_onnx_dynamo(self):
         mid = "arnir0/Tiny-LLM"
         summary, data = validate_model(
             mid,
             do_run=True,
             verbose=10,
             exporter="onnx-dynamo",
-            dump_folder="dump_test_validate_model_onnx",
+            dump_folder="dump_test_validate_model_onnx_dynamo",
             patch=True,
             stop_if_static=2 if pv.Version(torch.__version__) > pv.Version("2.6.1") else 0,
+            optimization="ir",
+        )
+        self.assertIsInstance(summary, dict)
+        self.assertIsInstance(data, dict)
+        self.assertLess(summary["disc_onnx_ort_run_abs"], 1e-4)
+        onnx_filename = data["onnx_filename"]
+        output_path = f"{onnx_filename}.ortopt.onnx"
+        run_ort_fusion(
+            onnx_filename, output_path, num_attention_heads=2, hidden_size=192, verbose=10
+        )
+
+    @requires_torch("2.7")
+    @hide_stdout()
+    @ignore_warnings(FutureWarning)
+    @requires_experimental()
+    def test_validate_model_custom(self):
+        mid = "arnir0/Tiny-LLM"
+        summary, data = validate_model(
+            mid,
+            do_run=True,
+            verbose=10,
+            exporter="custom",
+            dump_folder="dump_test_validate_model_custom",
+            patch=True,
+            stop_if_static=2 if pv.Version(torch.__version__) > pv.Version("2.6.1") else 0,
+            optimization="default",
         )
         self.assertIsInstance(summary, dict)
         self.assertIsInstance(data, dict)

onnx_diagnostic/ext_test_case.py

@@ -461,6 +461,25 @@ def requires_sklearn(version: str, msg: str = "") -> Callable:
     return lambda x: x
 
 
+def requires_experimental(version: str = "", msg: str = "") -> Callable:
+    """Skips a unit test if :epkg:`experimental-experiment` is not recent enough."""
+    import packaging.version as pv
+
+    try:
+        import experimental_experiment
+    except ImportError:
+        msg = f"experimental-experiment not installed: {msg}"
+        return unittest.skip(msg)
+
+    if pv.Version(experimental_experiment.__version__) < pv.Version(version):
+        msg = (
+            f"experimental-experiment version "
+            f"{experimental_experiment.__version__} < {version}: {msg}"
+        )
+        return unittest.skip(msg)
+    return lambda x: x
+
+
 def has_torch(version: str) -> bool:
     "Returns True if torch transformers is higher."
     import packaging.version as pv