onnx-diagnostic 0.8.0__py3-none-any.whl → 0.8.2__py3-none-any.whl

onnx_diagnostic/__init__.py

@@ -3,5 +3,5 @@ Patches, Investigates onnx models.
 Functions, classes to dig into a model when this one is right, slow, wrong...
 """
 
-__version__ = "0.8.0"
+__version__ = "0.8.2"
 __author__ = "Xavier Dupré"

onnx_diagnostic/_command_lines_parser.py

@@ -265,7 +265,7 @@ def get_parser_config() -> ArgumentParser:
         "--mop",
         metavar="KEY=VALUE",
         nargs="*",
-        help="Additional model options, use to change some parameters of the model, "
+        help="Additional model options, used to change some parameters of the model, "
         "example:\n  --mop attn_implementation=sdpa or --mop attn_implementation=eager",
         action=_ParseDict,
     )
@@ -442,11 +442,17 @@ def get_parser_validate(name: str = "validate") -> ArgumentParser:
         default=True,
         action=_BoolOrParseDictPatch,
         nargs="*",
-        help="Applies patches before exporting, it can be a boolean "
-        "to enable to disable the patches or be more finetuned. It is possible to "
-        "disable patch for torch by adding "
-        '--patch "patch_sympy=False" --patch "patch_torch=False", '
-        "default is True.",
+        help=textwrap.dedent(
+            """
+        Applies patches before exporting, it can be a boolean
+        to enable to disable the patches or be more finetuned
+        (default is True). It is possible to disable patch for torch
+        by adding:
+            --patch "patch_sympy=False" --patch "patch_torch=False"
+        """.strip(
+                "\n"
+            )
+        ),
     )
     parser.add_argument(
         "--rewrite",
@@ -476,10 +482,16 @@ def get_parser_validate(name: str = "validate") -> ArgumentParser:
         "--inputs2",
         default=1,
         type=int,
-        help="Validates or exports the model on a second set of inputs\n"
-        "to check the exported model supports dynamism. The values is used "
-        "as an increment to the first set of inputs. A high value may trick "
-        "a different behavior in the model and missed by the exporter.",
+        help=textwrap.dedent(
+            """
+        Validates or exports the model on a second set of inputs
+        to check the exported model supports dynamism. The values is used
+        as an increment to the first set of inputs. A high value may trick
+        a different behavior in the model and missed by the exporter.
+        """.strip(
+                "\n"
+            )
+        ),
     )
     parser.add_argument(
         "--runtime",
@@ -512,9 +524,15 @@ def get_parser_validate(name: str = "validate") -> ArgumentParser:
         parser.add_argument(
             "--ortfusiontype",
             required=False,
-            help="Applies onnxruntime fusion, this parameter should contain the\n"
-            "model type or multiple values separated by `|`. `ALL` can be used\n"
-            "to run them all.",
+            help=textwrap.dedent(
+                """
+                Applies onnxruntime fusion, this parameter should contain the
+                model type or multiple values separated by `|`. `ALL` can be used
+                to run them all.
+                """.strip(
+                    "\n"
+                )
+            ),
         )
     parser.add_argument("-v", "--verbose", default=0, type=int, help="verbosity")
     parser.add_argument("--dtype", help="Changes dtype if necessary.")
@@ -523,18 +541,32 @@ def get_parser_validate(name: str = "validate") -> ArgumentParser:
         "--iop",
         metavar="KEY=VALUE",
         nargs="*",
-        help="Additional input options, use to change the default"
-        "inputs use to export, example:\n  --iop cls_cache=SlidingWindowCache"
-        "\n  --iop cls_cache=StaticCache",
+        help=textwrap.dedent(
+            """
+        Additional input options, used to change the default
+        inputs use to export. Examples:
+            --iop cls_cache=SlidingWindowCache
+            --iop cls_cache=StaticCache
+        """.strip(
+                "\n"
+            )
+        ),
         action=_ParseDict,
     )
     parser.add_argument(
         "--mop",
         metavar="KEY=VALUE",
         nargs="*",
-        help="Additional model options, use to change some parameters of the model, "
-        "example:\n  --mop attn_implementation=sdpa --mop attn_implementation=eager\n  "
-        "--mop \"rope_scaling={'rope_type': 'dynamic', 'factor': 10.0}\"",
+        help=textwrap.dedent(
+            """
+            Additional model options, used to change some parameters
+            of the model. Example:
+                --mop attn_implementation=sdpa --mop attn_implementation=eager"
+                --mop "rope_scaling={'rope_type': 'dynamic', 'factor': 10.0}"
+            """.strip(
+                "\n"
+            )
+        ),
         action=_ParseDict,
     )
     if name == "validate":
@@ -566,9 +598,32 @@ def get_parser_validate(name: str = "validate") -> ArgumentParser:
         parser.add_argument(
             "--quiet-input-sets",
             default="",
-            help="Avoids raising an exception when an input sets does not work with "
-            "the exported model.\nExample: --quiet-input-sets=inputs,inputs22",
+            help=textwrap.dedent(
+                """
+                Avoids raising an exception when an input sets does not work with
+                the exported model. Example:
+                    --quiet-input-sets=inputs,inputs22
+                """.strip(
+                    "\n"
+                )
+            ),
         )
+    parser.add_argument(
+        "--expop",
+        metavar="KEY=VALUE",
+        nargs="*",
+        help=textwrap.dedent(
+            """
+            Additional exporter options, use to change some parameters
+            of the model. Examples:
+                --expop report=True
+                --expop report=True --expop verify=True
+            """.strip(
+                "\n"
+            )
+        ),
+        action=_ParseDict,
+    )
     return parser
 
 
@@ -634,6 +689,7 @@ def _cmd_validate(argv: List[Any]):
             output_names=(
                 None if len(args.outnames.strip()) < 2 else args.outnames.strip().split(",")
             ),
+            exporter_options=args.expop,
         )
         print("")
         print("-- summary --")
@@ -940,7 +996,7 @@ def get_parser_agg() -> ArgumentParser:
         "n_model_faster2x,n_model_faster3x,n_model_faster4x,n_node_attention,"
         "n_node_attention23,n_node_rotary_embedding,n_node_rotary_embedding23,"
         "n_node_gqa,n_node_layer_normalization,n_node_layer_normalization23,"
-        "peak_gpu_torch,peak_gpu_nvidia,n_node_control_flow,"
+        "peak_gpu_torch,peak_gpu_nvidia,n_node_control_flow,n_node_random,"
         "n_node_constant,n_node_shape,n_node_expand,"
         "n_node_function,n_node_initializer,n_node_scatter,"
         "time_export_unbiased,onnx_n_nodes_no_cst,n_node_initializer_small",

onnx_diagnostic/export/api.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Sequence, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
 import torch
 
 
@@ -14,6 +14,10 @@ def to_onnx(
     output_names: Optional[List[str]] = None,
     output_dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
     exporter: str = "onnx-dynamo",
+    exporter_kwargs: Optional[Dict[str, Any]] = None,
+    save_ep: Optional[str] = None,
+    optimize: bool = True,
+    use_control_flow_dispatcher: bool = False,
 ) -> Any:
     """
     Common API for exporters. By default, the models are optimized to use the
@@ -32,6 +36,11 @@ def to_onnx(
     :param output_names: to change the output of the onnx model
     :param output_dynamic_shapes: to overwrite the dynamic shapes names
     :param exporter: exporter to use (``onnx-dynamo``, ``modelbuilder``, ``custom``)
+    :param exporter_kwargs: additional parameters sent to the exporter
+    :param save_ep: saves the exported program
+    :param optimize: optimizes the model
+    :param use_control_flow_dispatcher: use the dispatcher created to supported
+        custom loops (see :func:`onnx_diagnostic.export.control_flow.loop_for`)
     :return: the output of the selected exporter, usually a structure including
         an onnx model
 
@@ -48,9 +57,23 @@ def to_onnx(
         )
     """
     if exporter == "custom":
-        from experimental_experiment.torch_interpreter import to_onnx as _to_onnx
+        from experimental_experiment.torch_interpreter import (
+            to_onnx as _to_onnx,
+            ExportOptions,
+        )
         from experimental_experiment.xbuilder import OptimizationOptions
 
+        if use_control_flow_dispatcher:
+            from .control_flow import create_global_dispatcher
+
+            dispatcher = create_global_dispatcher()
+
+        options = None
+        if exporter_kwargs is not None:
+            options = exporter_kwargs.pop("options", None)
+        if options is None:
+            options = OptimizationOptions(patterns="default+onnxruntime")
+
         return _to_onnx(
             mod,
             args=args,
@@ -63,7 +86,10 @@ def to_onnx(
             dynamic_shapes=dynamic_shapes,
             large_model=True,
             output_dynamic_shapes=output_dynamic_shapes,
-            options=OptimizationOptions(patterns="default+onnxruntime"),
+            export_options=ExportOptions(save_ep=save_ep),
+            options=options,
+            **(exporter_kwargs or {}),
+            dispatcher=dispatcher if use_control_flow_dispatcher else None,
         )
     if exporter in ("dynamo", "onnx-dynamo"):
         import onnxscript.rewriter.ort_fusions as ort_fusions
@@ -80,9 +106,12 @@ def to_onnx(
             opset_version=target_opset,
             dynamic_shapes=dynamic_shapes,
             dynamo=True,
+            **(exporter_kwargs or {}),
         )
-        ort_fusions.optimize_for_ort(epo.model)
-        epo.save(filename)
+        if optimize:
+            ort_fusions.optimize_for_ort(epo.model)
+        if filename:
+            epo.save(filename, external_data=True)
         return epo
 
     if exporter == "modelbuilder":
@@ -117,6 +146,7 @@ def to_onnx(
             precision=str(first_float[0].dtype).split(".")[-1],
             execution_provider="cuda" if first.is_cuda else "cpu",
             cache_dir=os.path.dirname(filename),
+            **(exporter_kwargs or {}),
         )
         save_model_builder(onx, os.path.dirname(filename))
         return onx