onnx-diagnostic 0.8.5__py3-none-any.whl → 0.8.7__py3-none-any.whl

onnx_diagnostic/ext_test_case.py

@@ -700,6 +700,19 @@ def requires_onnx(version: str, msg: str = "") -> Callable:
     return lambda x: x
 
 
+def requires_experimental_experiment(version: str, msg: str = "") -> Callable:
+    """Skips a unit test if :epkg:`onnx-array-api` is not recent enough."""
+    import packaging.version as pv
+    import experimental_experiment
+
+    if pv.Version(experimental_experiment.__version__) < pv.Version(version):
+        msg = (
+            f"onnx-array-api version {experimental_experiment.__version__} < {version}: {msg}"
+        )
+        return unittest.skip(msg)
+    return lambda x: x
+
+
 def requires_onnx_array_api(version: str, msg: str = "") -> Callable:
     """Skips a unit test if :epkg:`onnx-array-api` is not recent enough."""
     import packaging.version as pv
@@ -774,6 +787,7 @@ class ExtTestCase(unittest.TestCase):
     def setUpClass(cls):
         logger = logging.getLogger("onnxscript.optimizer.constant_folding")
         logger.setLevel(logging.ERROR)
+        warnings.filterwarnings("ignore", category=DeprecationWarning)
         unittest.TestCase.setUpClass()
 
     @classmethod
@@ -1253,6 +1267,7 @@ class ExtTestCase(unittest.TestCase):
             :class:`onnx_diagnostic.helpers.ort_session.InferenceSessionForTorch`
         """
         from .helpers import string_type, string_diff, max_diff
+        from .helpers.torch_helper import torch_deepcopy
         from .helpers.rt_helper import make_feeds
         from .helpers.ort_session import InferenceSessionForTorch
 
@@ -1269,6 +1284,12 @@ class ExtTestCase(unittest.TestCase):
                 model_file = proto
                 name = proto
                 proto = onnx.load(name)
+            elif hasattr(proto, "save"):
+                name = f"{test_name}.onnx"
+                proto.save(name)
+                proto = onnx.load(name)
+            elif hasattr(proto, "model_proto"):
+                proto = proto.model_proto
             elif not self.unit_test_going():
                 assert isinstance(
                     proto, onnx.ModelProto
@@ -1327,9 +1348,9 @@ class ExtTestCase(unittest.TestCase):
             if copy_inputs:
                 expected = [
                     (
-                        model(*copy.deepcopy(inp))
+                        model(*torch_deepcopy(inp))
                         if isinstance(inp, tuple)
-                        else model(**copy.deepcopy(inp))
+                        else model(**torch_deepcopy(inp))
                     )
                     for inp in inputs
                 ]

onnx_diagnostic/helpers/bench_run.py

@@ -20,7 +20,7 @@ class BenchmarkError(RuntimeError):
 
 
 def _clean_string(s: str) -> str:
-    cleaned = [c for c in s if 32 <= ord(c) < 127 and c not in {","}]
+    cleaned = [c for c in s if 32 <= ord(c) < 127 and c not in {",", ":"}]
     return "".join(cleaned)
 
 

onnx_diagnostic/helpers/log_helper.py

@@ -1921,9 +1921,7 @@ class CubeLogsPerformance(CubeLogs):
             return lambdas[formula]
 
         if formula == "onnx_n_nodes_no_cst":
-            return lambda df: gdf(df, "onnx_n_nodes", 0) - gdf(
-                df, "op_onnx__Constant", 0
-            ).fillna(0)
+            return lambda df: gdf(df, "onnx_n_nodes", 0) - gdf(df, "op_onnx__Constant", 0)
         if formula == "peak_gpu_torch":
             return lambda df: gdf(df, "mema_gpu_5_after_export") - gdf(df, "mema_gpu_4_reset")
         if formula == "peak_gpu_nvidia":

onnx_diagnostic/helpers/optim_helper.py

@@ -0,0 +1,116 @@
+from typing import Optional, Union
+import pprint
+import onnx
+
+
+def optimize_model(
+    algorithm: str,
+    model: Union[onnx.ModelProto, str],
+    output: Optional[str] = None,
+    processor: Optional[str] = None,
+    infer_shapes: bool = True,
+    remove_shape_info: bool = False,
+    verbose: int = 1,
+):
+    """
+    Optimizes an onnx model by fusing nodes. It looks for patterns in the graphs
+    and replaces them by the corresponding nodes. It also does basic optimization
+    such as removing identity nodes or unused nodes.
+
+    :param algorithm: algorithm to choose
+    :param model: model to optimize as a proto or a filename
+    :param output: if not empty, the optimized model is saved
+    :param processor: optimization are done for the processor
+    :param infer_shapes: infer shapes before optimizing, this might not be
+        available for all algorithm
+    :param remove_shape_info: remove shape information before saving the model
+    :param verbose: verbosity level
+    :return: optimized model
+
+    The goal is to make the model faster.
+    Argument patterns defines the patterns to apply or the set of patterns.
+    It is possible to show statistics or to remove a particular pattern.
+    Here are some environment variables which can be used to trigger
+    these displays.
+
+    Available options algorithms, default and default+runtime:
+
+    - ``DROPPATTERN=<pattern1,patterns2,...>``: do not apply
+      those patterns when optimizing a model
+    - ``DUMPPATTERNS=<folder>``: dumps all matched and applied nodes when a pattern is applied
+    - ``PATTERN=<pattern1,pattern2,...>``: increase verbosity
+      for specific patterns to understand why one pattern was not applied,
+      this shows which line is rejecting a pattern if it seems one pattern was missed
+    """
+    if isinstance(model, str):
+        if verbose:
+            print(f"[optimize_model] load {model!r}")
+        proto = onnx.load(model)
+        if verbose:
+            print("[optimize_model] done loading.")
+    else:
+        proto = model
+
+    if verbose:
+        print(f"[optimize_model] optimize with {algorithm!r}")
+    if algorithm in {"default", "default+onnxruntime"}:
+        from experimental_experiment.xoptim import get_pattern_list
+        from experimental_experiment.xbuilder import GraphBuilder, OptimizationOptions
+
+        pats = get_pattern_list(algorithm)
+
+        gr = GraphBuilder(
+            proto,
+            infer_shapes_options=infer_shapes,
+            optimization_options=OptimizationOptions(
+                patterns=pats,
+                verbose=verbose,
+                remove_unused=True,
+                constant_folding=True,
+                remove_identity=True,
+                max_iter=max(100, len(proto.graph.node) // 2),
+                processor=processor or "CPU",
+            ),
+        )
+        if verbose:
+            print(f"[optimize_model] starts optimizing with {len(pats)} patterns")
+            print(f"[optimize_model] model has {len(proto.graph.node)} nodes")
+        opt_onx, report = gr.to_onnx(optimize=True, return_optimize_report=True)
+        if verbose:
+            print("[optimize_model] optimization report")
+            pprint.pprint(report)
+            print("[optimize_model] done")
+
+    elif algorithm == "slim":
+        import onnxslim
+
+        opt_onx = onnxslim.slim(proto, no_shape_infer=not infer_shapes)
+    elif algorithm in {"ir", "os_ort"}:
+        import onnx_ir
+        import onnxscript.optimizer
+        from onnxscript.rewriter.ort_fusions import optimize_for_ort
+
+        model_ir = onnx_ir.from_proto(proto)
+        if algorithm == "ir":
+            onnxscript.optimizer.optimize(model_ir)
+        else:
+            optimize_for_ort(model_ir)
+        opt_onx = onnx_ir.serde.serialize_model(model_ir)
+
+    del proto
+    if verbose:
+        print(f"[optimize_model] done optimizing, model has {len(opt_onx.graph.node)} nodes")
+    if remove_shape_info:
+        if verbose:
+            print(f"[optimize_model] remove shape information {len(opt_onx.graph.value_info)}")
+        del opt_onx.graph.value_info[:]
+        if verbose:
+            print("[optimize_model] done removing shape info")
+
+    if output:
+        if verbose:
+            print(f"[optimize_model] save file into {output!r}")
+        onnx.save(opt_onx, output, save_as_external_data=True)
+        if verbose:
+            print("[optimize_model] done saving")
+    return opt_onx

onnx_diagnostic/tasks/image_text_to_text.py

@@ -13,6 +13,10 @@ from .data import get_data
 __TASK__ = "image-text-to-text"
 
 
+def should_have_vision_config(config):
+    return config.architectures != ["FuyuForCausalLM"]
+
+
 def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     kwargs: Dict[str, Any] = {}
@@ -477,7 +481,8 @@ def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]:
                 "hidden_size",
                 "pad_token_id",
             )
-            check_hasattr(config, "vision_config", ("image_token_index", "image_token_id"))
+            if should_have_vision_config(config):
+                check_hasattr(config, "vision_config", ("image_token_index", "image_token_id"))
             text_config = True
         else:
             check_hasattr(
@@ -491,7 +496,8 @@ def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]:
                 "vision_config",
             )
             text_config = False
-        check_hasattr(config.vision_config, ("num_channels", "in_chans", "in_channels"))
+        if should_have_vision_config(config):
+            check_hasattr(config.vision_config, ("num_channels", "in_chans", "in_channels"))
     kwargs = dict(
         head_dim=(
             16
@@ -552,17 +558,21 @@ def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]:
         ),
         width=(
             224
-            if config is None or not hasattr(config.vision_config, "image_size")
+            if config is None
+            or not should_have_vision_config(config)
+            or not hasattr(config.vision_config, "image_size")
             else config.vision_config.image_size
         ),
         height=(
             224
-            if config is None or not hasattr(config.vision_config, "image_size")
+            if config is None
+            or not should_have_vision_config(config)
+            or not hasattr(config.vision_config, "image_size")
             else config.vision_config.image_size
         ),
         num_channels=(
             3
-            if config is None
+            if config is None or not should_have_vision_config(config)
             else _pick(config.vision_config, "num_channels", "in_chans", "in_channels")
         ),
         pad_token_id=(

onnx_diagnostic/tasks/text2text_generation.py

@@ -18,6 +18,22 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
         config.num_decoder_layers = min(config.num_decoder_layers, 2)
     if hasattr(config, "num_hidden_layers"):
         config.num_hidden_layers = min(config.num_hidden_layers, nhl())
+    if hasattr(config, "encoder") and hasattr(config.encoder, "layer_types"):
+        default_layer_types = [
+            "sliding_attention",
+            "full_attention",
+            "sliding_attention",
+            "full_attention",
+        ]
+        config.encoder.num_hidden_layers = 4
+        config.encoder.layer_types = (
+            default_layer_types if config is None else config.encoder.layer_types[:4]
+        )
+        config.decoder.num_hidden_layers = 4
+        config.decoder.layer_types = (
+            default_layer_types if config is None else config.decoder.layer_types[:4]
+        )
+
     update_config(config, kwargs)
     return kwargs
 
@@ -177,55 +193,75 @@ def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]:
 
     If the configuration is None, the function selects typical dimensions.
     """
+    path = 1
     if config is not None:
-        check_hasattr(
-            config,
-            "vocab_size",
-            "hidden_size",
-            "num_attention_heads",
-            ("num_hidden_layers", "num_layers"),
-            ("n_positions", "d_model"),
-            (
-                "num_key_value_heads",
-                "num_heads",
-                ("decoder_attention_heads", "encoder_attention_heads"),
-            ),
-        )
-    # exceptions = {
-    #     "PLBartForConditionalGeneration": (
-    #         lambda c: c.encoder_attention_heads + c.decoder_attention_heads
-    #    )
-    # }
-    kwargs = dict(
-        batch_size=2,
-        sequence_length=30,
-        sequence_length2=3,
-        head_dim_encoder=16 if config is None else _pick(config, "d_kv", "encoder_ffn_dim"),
-        head_dim_decoder=16 if config is None else _pick(config, "d_kv", "decoder_ffn_dim"),
-        dummy_max_token_id=31999 if config is None else config.vocab_size - 1,
-        num_hidden_layers=(
-            8 if config is None else _pick(config, "num_hidden_layers", "num_layers")
-        ),
-        num_key_value_heads_encoder=(
-            16
-            if config is None
-            else _pick(
+        if hasattr(config, "num_attention_heads"):
+            check_hasattr(
                 config,
-                "encoder_attention_heads",
-                "num_key_value_heads",
-                "num_heads",
+                "vocab_size",
+                "hidden_size",
+                "num_attention_heads",
+                ("num_hidden_layers", "num_layers"),
+                ("n_positions", "d_model"),
+                (
+                    "num_key_value_heads",
+                    "num_heads",
+                    ("decoder_attention_heads", "encoder_attention_heads"),
+                ),
             )
-        ),
-        num_key_value_heads_decoder=(
-            16
-            if config is None
-            else _pick(
-                config,
-                "decoder_attention_heads",
-                "num_key_value_heads",
-                "num_heads",
-            )
-        ),
-        encoder_dim=512 if config is None else _pick(config, "n_positions", "d_model"),
-    )
+        else:
+            check_hasattr(config, "encoder", "decoder")
+            path = 2
+
+    if path == 1:
+        kwargs = dict(
+            batch_size=2,
+            sequence_length=30,
+            sequence_length2=3,
+            head_dim_encoder=(
+                16 if config is None else _pick(config, "d_kv", "encoder_ffn_dim")
+            ),
+            head_dim_decoder=(
+                16 if config is None else _pick(config, "d_kv", "decoder_ffn_dim")
+            ),
+            dummy_max_token_id=31999 if config is None else config.vocab_size - 1,
+            num_hidden_layers=(
+                8 if config is None else _pick(config, "num_hidden_layers", "num_layers")
+            ),
+            num_key_value_heads_encoder=(
+                16
+                if config is None
+                else _pick(
+                    config,
+                    "encoder_attention_heads",
+                    "num_key_value_heads",
+                    "num_heads",
+                )
+            ),
+            num_key_value_heads_decoder=(
+                16
+                if config is None
+                else _pick(
+                    config,
+                    "decoder_attention_heads",
+                    "num_key_value_heads",
+                    "num_heads",
+                )
+            ),
+            encoder_dim=512 if config is None else _pick(config, "n_positions", "d_model"),
+        )
+    else:
+        kwargs = dict(
+            batch_size=2,
+            sequence_length=30,
+            sequence_length2=3,
+            dummy_max_token_id=config.encoder.vocab_size - 1,
+            num_key_value_heads_encoder=config.encoder.num_key_value_heads,
+            num_key_value_heads_decoder=config.decoder.num_key_value_heads,
+            num_hidden_layers=len(config.encoder.layer_types),
+            head_dim_encoder=config.encoder.head_dim,
+            head_dim_decoder=config.decoder.head_dim,
+            encoder_dim=256,
+        )
+
     return kwargs, get_inputs

onnx_diagnostic/tasks/text_generation.py

@@ -40,6 +40,9 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
             state_size=8 if config is None else getattr(config, "state_size", None),
             conv_kernel=4 if config is None else getattr(config, "conv_kernel", None),
         )
+    elif config.__class__.__name__ == "FunnelConfig":
+        # does not support num_hidden_layers
+        kwargs = dict()
     else:
         kwargs = dict(
             head_dim=getattr(

onnx_diagnostic/torch_export_patches/onnx_export_errors.py

@@ -221,6 +221,7 @@ def _patch_torch(
     catch_constraints: bool,
     stop_if_static: int,
 ) -> Tuple[Optional[Callable], ...]:
+    import packaging.version as pv
     import torch
     import torch.jit
     import torch._export.non_strict_utils  # produce_guards_and_solve_constraints
@@ -238,6 +239,11 @@ def _patch_torch(
         patched_ShapeEnv,
     )
 
+    if pv.Version(torch.__version__) >= pv.Version("2.9.99"):
+        from .patches.patch_torch import patched_DynamicDimConstraintPrinter
+    else:
+        patched_DynamicDimConstraintPrinter = None
+
     f___constrain_user_specified_dimhint_range = None
     f__broadcast_in_dim_meta = None
     f__broadcast_shapes = None
@@ -259,6 +265,17 @@ def _patch_torch(
         print(f"[torch_export_patches] stop_if_static={stop_if_static!r}")
         print("[torch_export_patches] patch pytorch")
 
+    # torch.tx.experimental.symbolic_shapes.DynamicDimConstraintPrinter._print_Symbol
+    if patched_DynamicDimConstraintPrinter is not None:
+        f__print_symbol = (
+            torch.fx.experimental.symbolic_shapes.DynamicDimConstraintPrinter._print_Symbol
+        )
+        torch.fx.experimental.symbolic_shapes.DynamicDimConstraintPrinter._print_Symbol = (
+            patched_DynamicDimConstraintPrinter._print_Symbol
+        )
+    else:
+        f__print_symbol = None
+
     # torch.vmap
     f_vmap = torch.vmap
     torch.vmap = patched_vmap
@@ -392,6 +409,7 @@ def _patch_torch(
         f_shape_env__log_guard,
         f_shape_env__set_replacement,
         f_vmap,
+        f__print_symbol,
     )
 
 
@@ -416,6 +434,7 @@ def _unpatch_torch(
     f_shape_env__log_guard: Optional[Callable],
     f_shape_env__set_replacement: Optional[Callable],
     f_vmap: Optional[Callable],
+    f__print_symbol: Optional[Callable],
 ):
     import torch
     import torch.jit
@@ -423,6 +442,10 @@ def _unpatch_torch(
     from torch.fx.experimental.symbolic_shapes import ShapeEnv
 
     # this should disappear when torch.jit is removed
+    if f__print_symbol is not None:
+        torch.fx.experimental.symbolic_shapes.DynamicDimConstraintPrinter._print_Symbol = (
+            f__print_symbol
+        )
     torch.vmap = f_vmap
     torch.jit.isinstance = f_jit_isinstance
     torch._dynamo.mark_static_address = f_mark_static_address
@@ -818,6 +841,7 @@ def torch_export_patches(
     rewrite: Optional[List[Callable]] = None,
     dump_rewriting: Optional[str] = None,
     patch_details: Optional[PatchDetails] = None,
+    profile: Optional[str] = None,
 ) -> Callable:
     """
     Tries to bypass some situations :func:`torch.export.export` does not support.
@@ -847,9 +871,12 @@ def torch_export_patches(
         this is done by function :func:`transform_method
         <onnx_diagnostic.torch_export_patches.patch_module.transform_method>`,
         its documentation provides possible values
-    :param dump_rewriting: dumps rewriting information in file beginning with that prefix
-    :param patch_details: if specified, this class is used to stored every rewritten done.
+    :param dump_rewriting: dumps rewriting information in file beginning with that prefix,
+        this only applied on the automated rewritings
+    :param patch_details: if specified, this class is used to stored every applied rewriting.
     :param verbose: to show which patches is applied
+    :param profile: starts profiling whatever is called inside the context manager,
+        output the profiling into a text file
 
     The list of available patches.
 
@@ -989,6 +1016,7 @@ def torch_export_patches(
                 f_shape_env__log_guard,
                 f_shape_env__set_replacement,
                 f_vmap,
+                f__print_Symbol,
             ) = _patch_torch(
                 verbose, patch_details, patch_torch, catch_constraints, stop_if_static
             )
@@ -1017,10 +1045,23 @@ def torch_export_patches(
         if verbose:
             print("[torch_export_patches] done patching")
 
+        if profile:
+            from pyinstrument import Profiler
+
+            profiler = Profiler()
+            profiler.start()
+        else:
+            profiler = None
+
         try:
             yield fct_callable
         finally:
 
+            if profiler:
+                profiler.stop()
+                with open(profile, "w") as f:
+                    f.write(profiler.output_html())
+
             # unpatch
 
             if verbose:
@@ -1051,6 +1092,7 @@ def torch_export_patches(
                     f_shape_env__log_guard,
                     f_shape_env__set_replacement,
                     f_vmap,
+                    f__print_Symbol,
                 )
 
             if patch_transformers:

onnx_diagnostic/torch_export_patches/patch_expressions.py

@@ -101,7 +101,10 @@ def patched_selector(fct: Callable, patched_fct: Callable) -> Callable:
 
 
 def patched_float_arange(start, end, step):
-    """Patched arange when start, end, step are floats."""
+    """
+    Patched arange when start, end, step are floats.
+    This patch should not be needed after 2.10.
+    """
     if is_torchdynamo_exporting():
         return torch.ops.patched.float_arange(start, end, step)
     else:

onnx_diagnostic/torch_export_patches/patch_module.py

@@ -596,33 +596,41 @@ class RewriteControlFlow(ast.NodeTransformer):
                     elts=[
                         *[
                             ast.Call(
-                                ast.Attribute(
-                                    value=ast.Name(id="torch", ctx=ast.Load()),
-                                    attr="arange",
-                                    ctx=ast.Load(),
-                                ),
-                                args=[
-                                    ast.Subscript(
-                                        value=ast.Attribute(
-                                            value=ast.Name(id=v, ctx=ast.Load()),
-                                            attr="shape",
+                                func=ast.Attribute(
+                                    value=ast.Call(
+                                        ast.Attribute(
+                                            value=ast.Name(id="torch", ctx=ast.Load()),
+                                            attr="arange",
                                             ctx=ast.Load(),
                                         ),
-                                        slice=ast.Constant(value=0, ctx=ast.Load()),
+                                        args=[
+                                            ast.Subscript(
+                                                value=ast.Attribute(
+                                                    value=ast.Name(id=v, ctx=ast.Load()),
+                                                    attr="shape",
+                                                    ctx=ast.Load(),
+                                                ),
+                                                slice=ast.Constant(value=0, ctx=ast.Load()),
+                                                ctx=ast.Load(),
+                                            ),
+                                        ],
+                                        keywords=[
+                                            ast.keyword(
+                                                arg="dtype",
+                                                value=ast.Attribute(
+                                                    value=ast.Name(id="torch", ctx=ast.Load()),
+                                                    attr="int64",
+                                                    ctx=ast.Load(),
+                                                ),
+                                            )
+                                        ],
                                         ctx=ast.Load(),
                                     ),
-                                ],
-                                keywords=[
-                                    ast.keyword(
-                                        arg="dtype",
-                                        value=ast.Attribute(
-                                            value=ast.Name(id="torch", ctx=ast.Load()),
-                                            attr="int64",
-                                            ctx=ast.Load(),
-                                        ),
-                                    )
-                                ],
-                                ctx=ast.Load(),
+                                    attr="unsqueeze",
+                                    ctx=ast.Load(),
+                                ),
+                                args=[ast.Constant(value=1)],
+                                keywords=[],
                             )
                             for v in scan_shape_vars
                         ],