onnx-diagnostic 0.7.14__py3-none-any.whl → 0.7.15__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.7.14"
+__version__ = "0.7.15"
 __author__ = "Xavier Dupré"

onnx_diagnostic/_command_lines_parser.py

@@ -553,6 +553,12 @@ def get_parser_validate() -> ArgumentParser:
         action=BooleanOptionalAction,
         help="Enables onnxruntime logging when the session is created",
     )
+    parser.add_argument(
+        "--quiet-input-sets",
+        default="",
+        help="Avoids raising an exception when an input sets does not work with "
+        "the exported model, example: --quiet-input-sets=inputs,inputs22",
+    )
     return parser
 
 
@@ -614,6 +620,7 @@ def _cmd_validate(argv: List[Any]):
             warmup=args.warmup,
             inputs2=args.inputs2,
             ort_logs=args.ort_logs,
+            quiet_input_sets=set(args.quiet_input_sets.split(",")),
             output_names=(
                 None if len(args.outnames.strip()) < 2 else args.outnames.strip().split(",")
             ),
@@ -834,7 +841,7 @@ def get_parser_agg() -> ArgumentParser:
         "n_model_pass,n_model_faster,"
         "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_layer_normalization,n_node_layer_normalization23,"
+        "n_node_gqa,n_node_layer_normalization,n_node_layer_normalization23,"
         "peak_gpu_torch,peak_gpu_nvidia,n_node_control_flow,"
         "n_node_constant,n_node_shape,n_node_expand,"
         "n_node_function,n_node_initializer,n_node_scatter,"

onnx_diagnostic/helpers/cache_helper.py

@@ -108,7 +108,7 @@ def flatten_unflatten_for_dynamic_shapes(
 
 def is_cache_dynamic_registered(fast: bool = False) -> bool:
     """
-    Tells class :class:`transformers.cache_utils.DynamicCache` can be
+    Tells if class :class:`transformers.cache_utils.DynamicCache` can be
     serialized and deserialized. Only then, :func:`torch.export.export`
     can export a model.
 

onnx_diagnostic/helpers/config_helper.py

@@ -95,7 +95,8 @@ def config_class_from_architecture(arch: str, exc: bool = False) -> Optional[typ
     mod_name = cls.__module__
     mod = importlib.import_module(mod_name)
     source = inspect.getsource(mod)
-    reg = re.compile("config: ([A-Za-z0-9]+)")
+    # [^O] avoids capturing Optional[Something]
+    reg = re.compile("config: ([^O][A-Za-z0-9]+)")
     fall = reg.findall(source)
     if len(fall) == 0:
         assert not exc, (

onnx_diagnostic/helpers/rt_helper.py

@@ -3,8 +3,6 @@ import numpy as np
 import onnx
 import torch
 from .helper import string_type, flatten_object
-from .torch_helper import to_numpy
-from .cache_helper import is_cache_dynamic_registered
 
 
 def name_type_to_onnx_dtype(name: str) -> int:
@@ -49,7 +47,7 @@ def make_feeds(
     assert (
         not check_flatten
         or not all(isinstance(obj, torch.Tensor) for obj in flat)
-        or not is_cache_dynamic_registered(fast=True)
+        # or not is_cache_dynamic_registered(fast=True)
         or len(flat) == len(torch.utils._pytree.tree_flatten(inputs)[0])
     ), (
         f"Unexpected number of flattened objects, "
@@ -57,6 +55,8 @@ def make_feeds(
         f"{string_type(torch.utils._pytree.tree_flatten(inputs)[0], with_shape=True)}"
     )
     if use_numpy:
+        from .torch_helper import to_numpy
+
         flat = [to_numpy(t) if isinstance(t, torch.Tensor) else t for t in flat]
     names = (
         [i.name for i in proto.graph.input]

onnx_diagnostic/tasks/image_text_to_text.py

@@ -186,12 +186,13 @@ def _get_inputs_gemma3(
             f"total_sequence_length={total_sequence_length} != 860 "
             f"for model {model.__class__.__name__}"
         )
-        assert (
-            head_dim == 256
-        ), f"head_dim={head_dim} != 256 for model {model.__class__.__name__}"
+        assert head_dim in (
+            256,
+            32,
+        ), f"head_dim={head_dim} not in (32, 256) for model {model.__class__.__name__}"
         assert n_images == 1, f"n_images={n_images} != 1 for model {model.__class__.__name__}"
-        assert num_key_value_heads == 4, (
-            f"num_key_value_heads={num_key_value_heads} != 256 "
+        assert num_key_value_heads in (1, 4), (
+            f"num_key_value_heads={num_key_value_heads} not in (1, 4) "
             f"for this model {model.__class__.__name__}"
         )
 

onnx_diagnostic/tasks/text_generation.py

@@ -19,6 +19,9 @@ __TASK__ = "text-generation"
 def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     # FalconMambaConfig: use_mambapy
+    if hasattr(config, "text_config"):
+        # The model is probably of mixture of models used only for text.
+        config = config.text_config
     check_hasattr(
         config,
         ("head_dim", ("hidden_size", "num_attention_heads"), "use_mambapy"),
@@ -308,6 +311,9 @@ def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]:
 
     If the configuration is None, the function selects typical dimensions.
     """
+    if hasattr(config, "text_config"):
+        # The model is probably of mixture of models used only for text.
+        config = config.text_config
     if config is not None:
         check_hasattr(
             config,

onnx_diagnostic/torch_export_patches/onnx_export_errors.py

@@ -422,7 +422,7 @@ def torch_export_patches(
                 )
             )
 
-        if stop_if_static:
+        if patch_torch and stop_if_static:
             ShapeEnv._log_guard_remember = ShapeEnv._log_guard
 
             if verbose:

onnx_diagnostic/torch_export_patches/onnx_export_serialization.py

@@ -12,17 +12,26 @@ from transformers.cache_utils import (
     StaticCache,
 )
 
-try:
-    from transformers.models.mamba.modeling_mamba import MambaCache
-except ImportError:
-    from transformers.cache_utils import MambaCache
-
 from ..helpers import string_type
 from .serialization import _lower_name_with_
 
 PATCH_OF_PATCHES: Set[Any] = set()
 
 
+def get_mamba_cache_cls() -> type:
+    try:
+        from transformers.models.mamba.modeling_mamba import MambaCache
+
+        return MambaCache
+    except ImportError:
+        try:
+            from transformers.cache_utils import MambaCache
+
+            return MambaCache
+        except ImportError:
+            return None
+
+
 def register_class_serialization(
     cls,
     f_flatten: Callable,
@@ -203,13 +212,6 @@ def serialization_functions(
                 # f_check=make_dynamic_cache([(torch.rand((4, 4, 4)), torch.rand((4, 4, 4)))]),
                 verbose=verbose,
             ),
-            MambaCache: lambda verbose=verbose: register_class_serialization(
-                MambaCache,
-                flatten_mamba_cache,
-                unflatten_mamba_cache,
-                flatten_with_keys_mamba_cache,
-                verbose=verbose,
-            ),
             EncoderDecoderCache: lambda verbose=verbose: register_class_serialization(
                 EncoderDecoderCache,
                 flatten_encoder_decoder_cache,
@@ -232,6 +234,17 @@ def serialization_functions(
                 verbose=verbose,
             ),
         }
+        MambaCache = get_mamba_cache_cls()
+        if MambaCache:
+            transformers_classes[MambaCache] = (
+                lambda verbose=verbose: register_class_serialization(
+                    MambaCache,
+                    flatten_mamba_cache,
+                    unflatten_mamba_cache,
+                    flatten_with_keys_mamba_cache,
+                    verbose=verbose,
+                )
+            )
         classes.update(transformers_classes)
 
     if patch_diffusers:
@@ -287,7 +300,12 @@ def unregister_class_serialization(cls: type, verbose: int = 0):
 
 def unregister_cache_serialization(undo: Dict[str, bool], verbose: int = 0):
     """Undo all registrations."""
-    cls_ensemble = {MambaCache, DynamicCache, EncoderDecoderCache} | set(undo)
+    MambaCache = get_mamba_cache_cls()
+    cls_ensemble = (
+        {DynamicCache, EncoderDecoderCache}
+        | set(undo)
+        | ({MambaCache} if MambaCache else set())
+    )
     for cls in cls_ensemble:
         if undo.get(cls.__name__, False):
             unregister_class_serialization(cls, verbose)

onnx_diagnostic/torch_export_patches/patches/patch_torch.py

@@ -88,7 +88,7 @@ def patch__check_input_constraints_for_graph(
 
 def patched_infer_size(a, b):
     """Patches ``torch._subclasses.fake_impls.infer_size``."""
-    from torch.fx.experimental.symbolic_shapes import guard_size_oblivious
+    from torch.fx.experimental.symbolic_shapes import guard_or_false
 
     dimsA = len(a)
     dimsB = len(b)
@@ -113,19 +113,19 @@ def patched_infer_size(a, b):
         # were not the case, we'd need to write this using torch.sym_or() or
         # something like that).
         try:
-            b1 = guard_size_oblivious(sizeA == 1)
+            b1 = guard_or_false(sizeA == 1)
         except torch.fx.experimental.symbolic_shapes.GuardOnDataDependentSymNode:
             b1 = False
         try:
-            b2 = guard_size_oblivious(sizeB == 1)
+            b2 = guard_or_false(sizeB == 1)
         except torch.fx.experimental.symbolic_shapes.GuardOnDataDependentSymNode:
             b2 = False
         try:
-            b3 = guard_size_oblivious(sizeA == sizeB)
+            b3 = guard_or_false(sizeA == sizeB)
         except torch.fx.experimental.symbolic_shapes.GuardOnDataDependentSymNode:
             b3 = False
         if b1 or b2 or b3:
-            expandedSizes[i] = sizeB if guard_size_oblivious(sizeA == 1) else sizeA
+            expandedSizes[i] = sizeB if guard_or_false(sizeA == 1) else sizeA
         else:
             # PATCHED: generic case, the dimension is known, no need to assert
             expandedSizes[i] = torch.sym_max(sizeA, sizeB)
@@ -137,7 +137,6 @@ def patched__broadcast_shapes(*_shapes):
     from functools import reduce
     from torch._prims_common import IntLike
     from torch.fx.experimental.symbolic_shapes import (
-        guard_size_oblivious,
         guard_or_false,
         is_nested_int,
     )
@@ -174,13 +173,15 @@ def patched__broadcast_shapes(*_shapes):
                     continue
             # PATCHED: two cases, if == for sure, no broadcast,
             # otherwise maybe broadcast with max(dimensions)
-            if guard_size_oblivious(common_shape[idx] == 1):
+            if guard_or_false(common_shape[idx] != 1):
+                pass
+            elif guard_or_false(common_shape[idx] == 1) or guard_or_false(shape[idx] != 1):
                 if shape[idx] < 0:
                     raise ValueError(
                         "Attempting to broadcast a dimension with negative length!"
                     )
                 common_shape[idx] = shape[idx]
-            elif guard_size_oblivious(shape[idx] != 1):
+            else:
                 common_shape[idx] = torch.sym_max(common_shape[idx], shape[idx])
 
     return common_shape
@@ -360,6 +361,10 @@ class patched_ShapeEnv:
                 },
             )
 
+            for source in self.var_to_sources.get(a, []):
+                if user_tb:
+                    self.specialization_stacks[source] = user_tb
+
             # PATCHED: removed lines
             # if config.print_specializations:
             #    self.log.warning(
@@ -973,15 +978,22 @@ def patched__broadcast_in_dim_meta(
                     new_strides.append(a.stride()[original_idx])
                 else:
                     new_strides.append(0)
+            # PATCHED: disabled this check
+            elif guard_or_false(a.shape[original_idx] != 1):
+                new_strides.append(a.stride()[original_idx])
             else:
-                # PATCHED: disabled this check
-                # torch._check(
-                #    a.shape[original_idx] == shape[idx],
-                #    lambda idx=idx, original_idx=original_idx: (
-                #        f"non-broadcasting semantics require "
-                #        f"{a.shape[original_idx]} == {shape[idx]}"
-                #    ),
-                # )
+                torch._check(
+                    a.shape[original_idx] == shape[idx],
+                    lambda idx=idx, original_idx=original_idx: (
+                        f"non-broadcasting semantics require "
+                        f"{a.shape[original_idx]} == {shape[idx]}, "
+                        f"{guard_or_false(a.shape[idx] != 1)}, "
+                        f"guard_or_false(a.shape[idx] == 1)="
+                        f"{guard_or_false(a.shape[idx] == 1)}, "
+                        f"a.stride()={a.stride()}, idx={idx}, "
+                        f"original_idx={original_idx}"
+                    ),
+                )
                 new_strides.append(a.stride()[original_idx])
             original_idx = original_idx + 1
         else:

onnx_diagnostic/torch_export_patches/patches/patch_transformers.py

@@ -1019,6 +1019,26 @@ def patched__compute_dynamic_ntk_parameters(
     return inv_freq, attention_factor
 
 
+def _get_rope_init_fn(self, layer_type=None) -> Callable:
+    if hasattr(self, "rope_init_fn"):
+        # transformers<=5.0
+        rope_init_fn = (
+            patched__compute_dynamic_ntk_parameters
+            if self.rope_init_fn
+            is transformers.modeling_rope_utils._compute_dynamic_ntk_parameters
+            else self.rope_init_fn
+        )
+        return rope_init_fn
+
+    rope_type = self.rope_type if layer_type is None else self.rope_type[layer_type]
+    rope_init_fn = self.compute_default_rope_parameters
+    if rope_type != "default":
+        rope_init_fn = transformers.modeling_rope_utils.ROPE_INIT_FUNCTIONS[self.rope_type]
+    if rope_init_fn is transformers.modeling_rope_utils._compute_dynamic_ntk_parameters:
+        return patched__compute_dynamic_ntk_parameters
+    return rope_init_fn
+
+
 def patched_dynamic_rope_update(rope_forward):
     """manual patch: ``[patch:transformers.modeling_rope_utils.dynamic_rope_update]``
 
@@ -1082,22 +1102,27 @@ def patched_dynamic_rope_update(rope_forward):
 
     """
 
-    def longrope_frequency_update(self, position_ids, device):
+    def longrope_frequency_update(self, position_ids, device, layer_type=None):
         # It is no use to patch the function after the model is created
         # as rope_init_fn is an attribute set to one function when the model
         # is created and when no patch is applied yet.
         # So we select the patched version here.
-        rope_init_fn = (
-            patched__compute_dynamic_ntk_parameters
-            if self.rope_init_fn
-            is transformers.modeling_rope_utils._compute_dynamic_ntk_parameters
-            else self.rope_init_fn
-        )
+        rope_init_fn = _get_rope_init_fn(self, layer_type=layer_type)
         seq_len = torch.max(position_ids) + 1
         if hasattr(self.config, "original_max_position_embeddings"):
             original_max_position_embeddings = self.config.original_max_position_embeddings
         else:
             original_max_position_embeddings = self.config.max_position_embeddings
+
+        if layer_type is None:
+            # rope_type = self.rope_type
+            original_inv_freq = self.original_inv_freq
+            prefix = ""
+        else:
+            # rope_type = self.rope_type[layer_type]
+            original_inv_freq = getattr(self, f"{layer_type}_original_inv_freq")
+            prefix = f"{layer_type}_"
+
         # At export time, seq_len is unknown.
         long_inv_freq, _ = rope_init_fn(
             self.config, device, seq_len=original_max_position_embeddings + 1
@@ -1112,13 +1137,13 @@ def patched_dynamic_rope_update(rope_forward):
             (lambda x, y: y.clone()),
             [long_inv_freq, original_inv_freq],
         )
-        self.inv_freq = inv_freq
+        setattr(self, f"{prefix}inv_freq", inv_freq)
         # if seq_len > original_max_position_embeddings:
         #    self.inv_freq = self.long_inv_freq
         # else:
         #    self.inv_freq = self.original_inv_freq
 
-    def dynamic_frequency_update(self, position_ids, device):
+    def dynamic_frequency_update(self, position_ids, device, layer_type=None):
         # constructor:
         # - self.max_seq_len_cached = config.max_position_embeddings
         # - self.original_max_seq_len = config.max_position_embeddings
@@ -1128,12 +1153,7 @@ def patched_dynamic_rope_update(rope_forward):
         # as rope_init_fn is an attribute set to one function when the model
         # is created and when no patch is applied yet.
         # So we select the patched version here.
-        rope_init_fn = (
-            patched__compute_dynamic_ntk_parameters
-            if self.rope_init_fn
-            is transformers.modeling_rope_utils._compute_dynamic_ntk_parameters
-            else self.rope_init_fn
-        )
+        rope_init_fn = _get_rope_init_fn(self, layer_type=layer_type)
 
         # This behaviour is difficult to translate.
         # The sequence always grows.
@@ -1162,6 +1182,19 @@ def patched_dynamic_rope_update(rope_forward):
             self.config, device, seq_len=seq_len
         )
 
+        if layer_type is None:
+            # rope_type = self.rope_type
+            # max_seq_len_cached = self.max_seq_len_cached
+            original_inv_freq = self.original_inv_freq
+            prefix = ""
+        else:
+            # rope_type = self.rope_type[layer_type]
+            # max_seq_len_cached = getattr(
+            #     self, f"{layer_type}_max_seq_len_cached", self.max_seq_len_cached
+            # )
+            original_inv_freq = getattr(self, f"{layer_type}_original_inv_freq")
+            prefix = f"{layer_type}_"
+
         # Second test to translate.
         # Let's keep in mind, self.max_seq_len_cached = seq_len is likely to be True.
         # But in that case the following condition is a way to restore the original cache.
@@ -1183,15 +1216,26 @@ def patched_dynamic_rope_update(rope_forward):
             (lambda x, y: y.clone()),
             [long_inv_freq, original_inv_freq],
         )
-        self.inv_freq = inv_freq
+        setattr(self, f"{prefix}inv_freq", inv_freq)
 
     @wraps(rope_forward)
-    def wrapper(self, x, position_ids):
+    def wrapper(self, x, position_ids, layer_type=None):
+        if layer_type is None:
+            if "dynamic" in self.rope_type:
+                dynamic_frequency_update(self, position_ids, device=x.device)
+            elif self.rope_type == "longrope":
+                longrope_frequency_update(self, position_ids, device=x.device)
+            return rope_forward(self, x, position_ids)
+
         if "dynamic" in self.rope_type:
-            dynamic_frequency_update(self, position_ids, device=x.device)
+            dynamic_frequency_update(
+                self, position_ids, device=x.device, layer_type=layer_type
+            )
         elif self.rope_type == "longrope":
-            longrope_frequency_update(self, position_ids, device=x.device)
-        return rope_forward(self, x, position_ids)
+            longrope_frequency_update(
+                self, position_ids, device=x.device, layer_type=layer_type
+            )
+        return rope_forward(self, x, position_ids, layer_type=layer_type)
 
     return wrapper
 
@@ -1287,12 +1331,18 @@ class common_RotaryEmbedding(torch.nn.Module):
     # @torch.no_grad()
     # PATCHED: the decorator
     @patched_dynamic_rope_update
-    def forward(self, x, position_ids):
+    def forward(self, x, position_ids, layer_type=None):
+        if layer_type is not None:
+            # transformers>=5.0
+            inv_freq = getattr(self, f"{layer_type}_inv_freq")
+            attention_scaling = getattr(self, f"{layer_type}_attention_scaling")
+        else:
+            # transformers<5.0
+            inv_freq = self.inv_freq
+            attention_scaling = self.attention_scaling
+
         inv_freq_expanded = (
-            self.inv_freq[None, :, None]
-            .float()
-            .expand(position_ids.shape[0], -1, 1)
-            .to(x.device)
+            inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
         )
         position_ids_expanded = position_ids[:, None, :].float()
 
@@ -1304,8 +1354,8 @@ class common_RotaryEmbedding(torch.nn.Module):
         with torch.autocast(device_type=device_type, enabled=False):  # Force float32
             freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
             emb = torch.cat((freqs, freqs), dim=-1)
-            cos = emb.cos() * self.attention_scaling
-            sin = emb.sin() * self.attention_scaling
+            cos = emb.cos() * attention_scaling
+            sin = emb.sin() * attention_scaling
 
         return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
 
@@ -1380,7 +1430,8 @@ class patched_IdeficsEmbedding(torch.nn.Module):
 
         def _set_cos_sin_cache_then(x, inv_freq, seq_len, _cos_cached, _sin_cached):
             t = torch.arange(seq_len, device=x.device, dtype=torch.int64).type_as(inv_freq)
-            freqs = torch.einsum("i,j->ij", t, inv_freq)
+            # freqs = torch.einsum("i,j->ij", t, inv_freq)
+            freqs = t.reshape((-1, 1)) * inv_freq.reshape((1, -1))
             emb = torch.cat((freqs, freqs), dim=-1)
             return emb.cos().to(x.dtype), emb.sin().to(x.dtype)
 

onnx_diagnostic/torch_models/hghub/model_inputs.py

@@ -95,6 +95,8 @@ def get_untrained_model_with_inputs(
         print("-- dynamic shapes:", pprint.pformat(data['dynamic_shapes']))
         print("-- configuration:", pprint.pformat(data['configuration']))
     """
+    if task == "":
+        task = None
     assert not use_preinstalled or not use_only_preinstalled, (
         f"model_id={model_id!r}, preinstalled model is only available "
         f"if use_only_preinstalled is False."
@@ -120,14 +122,16 @@ def get_untrained_model_with_inputs(
             **(model_kwargs or {}),
         )
 
-    model, task, mkwargs, diff_config = None, None, {}, None
+    model, task_, mkwargs, diff_config = None, None, {}, None
     if use_pretrained and same_as_pretrained:
         if model_id in HANDLED_MODELS:
-            model, task, config = load_specific_model(model_id, verbose=verbose)
+            model, task_, config = load_specific_model(model_id, verbose=verbose)
 
+    if task is None:
+        task = task_
     if model is None:
         arch = architecture_from_config(config)
-        if arch is None:
+        if task is None and arch is None:
             task = task_from_id(model_id, subfolder=subfolder)
         assert task is not None or arch is not None, (
             f"Unable to determine the architecture for model {model_id!r}, "

onnx_diagnostic/torch_models/validate.py

@@ -4,7 +4,7 @@ import inspect
 import os
 import pprint
 import sys
-from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Sequence, Set, Tuple, Union
 import time
 import numpy as np
 import onnx
@@ -117,11 +117,21 @@ def _make_folder_name(
     drop_inputs: Optional[List[str]] = None,
     same_as_pretrained: bool = False,
     use_pretrained: bool = False,
+    task: Optional[str] = None,
 ) -> str:
     "Creates a filename unique based on the given options."
     els = [model_id.replace("/", "_")]
     if subfolder:
         els.append(subfolder.replace("/", "_"))
+    if not task:
+        els.append(task)  # type: ignore[arg-type]
+    if drop_inputs:
+        ii = "-".join(f"{s[0]}{s[-1]}" for s in drop_inputs)
+        els.append(f"I-{ii.upper()}")
+    if use_pretrained:
+        els.append("TRAINED")
+    elif same_as_pretrained:
+        els.append("SAMESIZE")
     if exporter:
         els.append(exporter)
     if optimization:
@@ -142,14 +152,7 @@ def _make_folder_name(
         els.append(sdev)
     if opset is not None:
         els.append(f"op{opset}")
-    if drop_inputs:
-        ii = "-".join(f"{s[0]}{s[-1]}" for s in drop_inputs)
-        els.append(f"I-{ii.upper()}")
-    if use_pretrained:
-        els.append("TRAINED")
-    elif same_as_pretrained:
-        els.append("SAMESIZE")
-    return "-".join(els)
+    return "/".join([e for e in els if e])
 
 
 def version_summary() -> Dict[str, Union[int, float, str]]:
@@ -319,6 +322,7 @@ def validate_model(
     inputs2: int = 1,
     output_names: Optional[List[str]] = None,
     ort_logs: bool = False,
+    quiet_input_sets: Optional[Set[str]] = None,
 ) -> Tuple[Dict[str, Union[int, float, str]], Dict[str, Any]]:
     """
     Validates a model.
@@ -373,6 +377,8 @@ def validate_model(
         or an empty cache for example
     :param output_names: output names the onnx exporter should use
     :param ort_logs: increases onnxruntime verbosity when creating the session
+    :param quiet_input_sets: avoid raising an exception if the inputs belongs to that set
+        even if quiet is False
     :return: two dictionaries, one with some metrics,
         another one with whatever the function produces
 
@@ -473,6 +479,7 @@ def validate_model(
             drop_inputs=drop_inputs,
             use_pretrained=use_pretrained,
             same_as_pretrained=same_as_pretrained,
+            task=task,
         )
         dump_folder = os.path.join(dump_folder, folder_name)
         if not os.path.exists(dump_folder):
@@ -487,6 +494,8 @@ def validate_model(
             print(f"[validate_model] validate model id {model_id!r}, subfolder={subfolder!r}")
         else:
             print(f"[validate_model] validate model id {model_id!r}")
+        if task:
+            print(f"[validate_model] with task {task!r}")
         print(f"[validate_model] patch={patch!r}")
         if model_options:
             print(f"[validate_model] model_options={model_options!r}")
@@ -762,6 +771,10 @@ def validate_model(
             ep = data["exported_program"]
             if verbose:
                 print(f"[validate_model] -- dumps exported program in {dump_folder!r}...")
+            assert isinstance(
+                folder_name, str
+            ), f"folder_name={folder_name!r} should be a string"
+            folder_name = folder_name.replace("/", "-")
             with open(os.path.join(dump_folder, f"{folder_name}.ep"), "w") as f:
                 f.write(str(ep))
             torch.export.save(ep, os.path.join(dump_folder, f"{folder_name}.pt2"))
@@ -770,6 +783,10 @@ def validate_model(
             if verbose:
                 print("[validate_model] done (dump ep)")
         if "onnx_program" in data:
+            assert isinstance(
+                folder_name, str
+            ), f"folder_name={folder_name!r} should be a string"
+            folder_name = folder_name.replace("/", "-")
             epo = data["onnx_program"]
             if verbose:
                 print(f"[validate_model] dumps onnx program in {dump_folder!r}...")
@@ -842,6 +859,7 @@ def validate_model(
             warmup=warmup,
             second_input_keys=second_input_keys,
             ort_logs=ort_logs,
+            quiet_input_sets=quiet_input_sets,
         )
         summary.update(summary_valid)
         summary["time_total_validation_onnx"] = time.perf_counter() - validation_begin
@@ -904,6 +922,7 @@ def validate_model(
                     repeat=repeat,
                     warmup=warmup,
                     second_input_keys=second_input_keys,
+                    quiet_input_sets=quiet_input_sets,
                 )
                 summary.update(summary_valid)
 
@@ -1289,6 +1308,7 @@ def validate_onnx_model(
     warmup: int = 0,
     second_input_keys: Optional[List[str]] = None,
     ort_logs: bool = False,
+    quiet_input_sets: Optional[Set[str]] = None,
 ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
     """
     Verifies that an onnx model produces the same
@@ -1308,6 +1328,7 @@ def validate_onnx_model(
         to make sure the exported model supports dynamism, the value is
         used as an increment added to the first set of inputs (added to dimensions)
     :param ort_logs: triggers the logs for onnxruntime
+    :param quiet_input_sets: avoid raising an exception for these sets of inputs
     :return: two dictionaries, one with some metrics,
         another one with whatever the function produces
     """
@@ -1431,6 +1452,8 @@ def validate_onnx_model(
     keys = [("inputs", "run_expected", "")]
     if second_input_keys:
         keys.extend([(k, f"run_expected2{k[6:]}", f"2{k[6:]}") for k in second_input_keys])
+    if verbose:
+        print(f"[validate_onnx_model] -- keys={keys}")
     for k_input, k_expected, suffix in keys:
         # make_feeds
         assert k_input in data, f"Unable to find {k_input!r} in {sorted(data)}"
@@ -1455,10 +1478,12 @@ def validate_onnx_model(
 
         # run ort
         if verbose:
-            print("[validate_onnx_model] run session...")
+            print(f"[validate_onnx_model] run session on inputs 'inputs{suffix}'...")
+            if quiet_input_sets and f"inputs{suffix}" in quiet_input_sets:
+                print(f"[validate_onnx_model] quiet_input_sets={quiet_input_sets}")
 
         got = _quiet_or_not_quiet(
-            quiet,
+            quiet or (quiet_input_sets is not None and f"inputs{suffix}" in quiet_input_sets),
             _mk(f"run_onnx_ort{suffix}"),
             summary,
             data,

{onnx_diagnostic-0.7.14.dist-info → onnx_diagnostic-0.7.15.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: onnx-diagnostic
-Version: 0.7.14
+Version: 0.7.15
 Summary: Tools to help converting pytorch models into ONNX.
 Home-page: https://github.com/sdpython/onnx-diagnostic
 Author: Xavier Dupré

{onnx_diagnostic-0.7.14.dist-info → onnx_diagnostic-0.7.15.dist-info}/RECORD

@@ -1,6 +1,6 @@
-onnx_diagnostic/__init__.py,sha256=fxgnYe-ZeX2ZhqiqehQfAUIDhdiy2BjpbzcaUtrI5g8,174
+onnx_diagnostic/__init__.py,sha256=rsSRl1QPe3XLEW8cPAdFVV6eLYgcRFg37pNNVTw5FAc,174
 onnx_diagnostic/__main__.py,sha256=YmyV_Aq_ianDlHyKLHMa6h8YK3ZmFPpLVHLKjM91aCk,79
-onnx_diagnostic/_command_lines_parser.py,sha256=qCPdI1_Za7OM1MuR1utyhTcSZQlM4UVmN8Su4HoRjvI,33670
+onnx_diagnostic/_command_lines_parser.py,sha256=bl6lorcodFk14dm2lgeCjr4OkRziUrKXn3GGkbxIuVA,33987
 onnx_diagnostic/api.py,sha256=BhCl_yCd78N7TlVtPOHjeYv1QBEy39TjZ647rcHqLh0,345
 onnx_diagnostic/doc.py,sha256=t3RELgfooYnVMAi0JSpggWkQEgUsREz8NmRvn0TnLI8,2829
 onnx_diagnostic/ext_test_case.py,sha256=emfQGiQSz5FVDhyJ1Acsv_Tast7tWl426TjtpNqxDBU,43558
@@ -12,8 +12,8 @@ onnx_diagnostic/helpers/__init__.py,sha256=GJ2GT7cgnlIveVUwMZhuvUwidbTJaKv8CsSIO
 onnx_diagnostic/helpers/_log_helper.py,sha256=OTwQH0OIxs9B6nrSvR7MoxMimSw_8mU0mj133NvLk5o,16832
 onnx_diagnostic/helpers/args_helper.py,sha256=SRWnqC7EENg09RZlA50B_PcdiIhdbgA4C3ACfzl5nMs,4419
 onnx_diagnostic/helpers/bench_run.py,sha256=CGA6VMJZMH2gDhVueT9ypNm4PMcjGrrGFYp08nhWj9k,16539
-onnx_diagnostic/helpers/cache_helper.py,sha256=zxjm0-3lHs0A7wLEejz2r2KPMPjkkva--8511MaSy74,24846
-onnx_diagnostic/helpers/config_helper.py,sha256=H2mOcMXfrcolFnt8EuqmRFkpQ3YdNRDfvm9ToI1vNH0,5618
+onnx_diagnostic/helpers/cache_helper.py,sha256=4fkPKLG590l1Gbqaw_QubXeJqY17X2Z6CBmLp23-BYI,24849
+onnx_diagnostic/helpers/config_helper.py,sha256=cWRETgFhZ7tayIZPnMqF8BF5AvTU64G2BMqyzgO7lzs,5670
 onnx_diagnostic/helpers/doc_helper.py,sha256=pl5MZd3_FaE8BqQnqoBuSBxoNCFcd2OJd3eITUSku5c,5897
 onnx_diagnostic/helpers/graph_helper.py,sha256=hevQT5a7_QuriVPQcbT5qe18n99Doyl5h3-qshx1-uk,14093
 onnx_diagnostic/helpers/helper.py,sha256=zl7vG6G4ueq931Z9iT8OlKfmtFxvRJD2WJQh_qsMiBs,63401
@@ -23,7 +23,7 @@ onnx_diagnostic/helpers/mini_onnx_builder.py,sha256=Cgx1ojmV0S_JpZ_UqwsNxeULMMDv
 onnx_diagnostic/helpers/model_builder_helper.py,sha256=sK40KRN9GWK1vbNJHIXkYAojblbKD0bdom7BFmoNSv4,12860
 onnx_diagnostic/helpers/onnx_helper.py,sha256=oxl3x0EQowGP9kfz8aKDqnJZcvYY8FeZLsfoLJDiSUg,39826
 onnx_diagnostic/helpers/ort_session.py,sha256=UgUUeUslDxEFBc6w6f3HMq_a7bn4TBlItmojqWquSj4,29281
-onnx_diagnostic/helpers/rt_helper.py,sha256=JnqsidpmX47ux5jaA_7Of_eS7KIRlOTqqDKo7ZUD-bI,5251
+onnx_diagnostic/helpers/rt_helper.py,sha256=mmxQ0RQ7mhG0ybHOtzbZiV2mt503JVaKcErQQ79ydWs,5208
 onnx_diagnostic/helpers/torch_helper.py,sha256=SY01uEx5tKtPcix91AifhgmsvNkDMGpTigT7w_0Nj98,34442
 onnx_diagnostic/reference/__init__.py,sha256=rLZsxOlnb7-81F2CzepGnZLejaROg4JvgFaGR9FwVQA,208
 onnx_diagnostic/reference/evaluator.py,sha256=RzNzjFDeMe-4X51Tb22N6aagazY5ktNq-mRmPcfY5EU,8848
@@ -77,7 +77,7 @@ onnx_diagnostic/tasks/automatic_speech_recognition.py,sha256=umZmjGW1gDUFkqvBJnQ
 onnx_diagnostic/tasks/feature_extraction.py,sha256=Zh9p_Q8FqEO2_aqI0cCiq8OXuM3WUZbwItlLOmLnNl8,5537
 onnx_diagnostic/tasks/fill_mask.py,sha256=5Gt6zlj0p6vuifox7Wmj-TpHXJvPS0CEH8evgdBHDNA,2640
 onnx_diagnostic/tasks/image_classification.py,sha256=nLpBBB1Gkog3Fk6pu2waiHcuQr4ILPptc9FhQ-pn460,4682
-onnx_diagnostic/tasks/image_text_to_text.py,sha256=EcaIdSYfaGLomSuO6G39lNd70tqFb19Xx3CjpQxQp9o,21538
+onnx_diagnostic/tasks/image_text_to_text.py,sha256=HDXuk1bEE3qTR0mUR_6rw-5RAXSyUvGY-dMNamIpvn0,21577
 onnx_diagnostic/tasks/image_to_video.py,sha256=SoF2cVIJr6P30Abp-FCuixFDh5RvTuNEOL36QthGY6U,3860
 onnx_diagnostic/tasks/mask_generation.py,sha256=fjdD3rd-O-mFL0hQy3la3JXKth_0bH2HL7Eelq-3Dbs,5057
 onnx_diagnostic/tasks/mixture_of_expert.py,sha256=al4tk1BrHidtRiHlAaiflWiJaAte0d5M8WcBioANG9k,2808
@@ -86,14 +86,14 @@ onnx_diagnostic/tasks/sentence_similarity.py,sha256=vPqNZgAnIvY0rKWPUTs0IlU3RFQD
 onnx_diagnostic/tasks/summarization.py,sha256=8vB_JiRzDEacIvr8CYTuVQTH73xG_jNkndoS9RHJTSs,8292
 onnx_diagnostic/tasks/text2text_generation.py,sha256=35eF_RlSeMdLTZPooLMAnszs-z0bkKZ34Iej3JgA96A,8602
 onnx_diagnostic/tasks/text_classification.py,sha256=CGc72SpXFzTUyzAHEMPgyy_s187DaYGsRdrosxG80_Q,2711
-onnx_diagnostic/tasks/text_generation.py,sha256=-oWq_I1lAUm9wxJnvFM1kXDJAmHbCiM6lUG3waR3o2k,13909
+onnx_diagnostic/tasks/text_generation.py,sha256=FwpmI4c_cO9uYQwJFfsHRMArPdwaeU5TBan2lisoHZk,14205
 onnx_diagnostic/tasks/text_to_image.py,sha256=mOS3Ruosi3hzRMxXLDN7ZkAbi7NnQb7MWwQP_okGVHs,2962
 onnx_diagnostic/tasks/zero_shot_image_classification.py,sha256=jJCMWuOqGv5ahCfjrcqxuYCJFhTgHV5KUf2yyv2yxYA,4624
 onnx_diagnostic/tasks/data/__init__.py,sha256=uJoemrWgEjI6oA-tMX7r3__x-b3siPmkgqaY7bgIles,401
 onnx_diagnostic/tasks/data/dummies_imagetext2text_generation_gemma3.onnx,sha256=UbtvmWMqcZOKJ-I-HXWI1A6YR6QDaFS5u_yXm5C3ZBw,10299
 onnx_diagnostic/torch_export_patches/__init__.py,sha256=0SaZedwznm1hQUCvXZsGZORV5vby954wEExr5faepGg,720
-onnx_diagnostic/torch_export_patches/onnx_export_errors.py,sha256=ZMsUeU3Hx5YD8xNgQTaW8Br88HvPSiCmqmKLhMz5jw0,30459
-onnx_diagnostic/torch_export_patches/onnx_export_serialization.py,sha256=klvqiMjccwGhiRnLRVbwTi5WWkMfvtnOV5ycirPcAdA,11354
+onnx_diagnostic/torch_export_patches/onnx_export_errors.py,sha256=T2FaIBSU3NfUyt54whwBmRHPuAzmZKFVHuwu-mikNz4,30475
+onnx_diagnostic/torch_export_patches/onnx_export_serialization.py,sha256=K78uX5EHTuu0ah3mkZWNcGow4775GKH-EnDs3ZlIEhE,11778
 onnx_diagnostic/torch_export_patches/patch_expressions.py,sha256=vr4tt61cbDnaaaduzMj4UBZ8OUtr6GfDpIWwOYqjWzs,3213
 onnx_diagnostic/torch_export_patches/patch_inputs.py,sha256=2HQZKQV6TM5430RIvKiMPe4cfGvFdx1UnP1w76CeGE4,8110
 onnx_diagnostic/torch_export_patches/patch_module.py,sha256=R2d9IHM-RwsBKDsxuBIJnEqMoxbS9gd4YWFGG2wwV5A,39881
@@ -101,19 +101,19 @@ onnx_diagnostic/torch_export_patches/patch_module_helper.py,sha256=2U0AdyZuU0W54
 onnx_diagnostic/torch_export_patches/eval/__init__.py,sha256=YQoOGt9XQLWqnJ15NnT7ri_jDevfvpuQwEJo38E-VRU,25056
 onnx_diagnostic/torch_export_patches/eval/model_cases.py,sha256=joDJV1YfrhYBR_6eXYvNO1jbiJM8Whb47NWZxo8SBwg,27172
 onnx_diagnostic/torch_export_patches/patches/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-onnx_diagnostic/torch_export_patches/patches/patch_torch.py,sha256=muA2i6Krd6iB2-nIteplxo_pvQEx4LQMZTxDmLe1n44,40825
-onnx_diagnostic/torch_export_patches/patches/patch_transformers.py,sha256=hHO7XOaUw3XrhPSrG2hTpMNzGVm_zigLg8d7hMOK7Gs,79188
+onnx_diagnostic/torch_export_patches/patches/patch_torch.py,sha256=QIm3GabPnoJDIM1HJl0reyUKf7fg7h57TsHkWfDWjF4,41408
+onnx_diagnostic/torch_export_patches/patches/patch_transformers.py,sha256=aVYEhrn48YUGn0rim5o2oygWFkwm3-HsGRpS1rGySeQ,81496
 onnx_diagnostic/torch_export_patches/serialization/__init__.py,sha256=BHLdRPtNAtNPAS-bPKEj3-foGSPvwAbZXrHzGGPDLEw,1876
 onnx_diagnostic/torch_export_patches/serialization/diffusers_impl.py,sha256=drq3EH_yjcSuIWYsVeUWm8Cx6YCZFU6bP_1PLtPfY5I,945
 onnx_diagnostic/torch_export_patches/serialization/transformers_impl.py,sha256=mcmZGekzQlLgE_o3SdKlRgCx4ewwyyAuNWZ9CaN_zrI,9317
 onnx_diagnostic/torch_models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 onnx_diagnostic/torch_models/llms.py,sha256=soyg4yC87ptGoeulJhKqw5opGmuLvH1pn_ZDXZ4Jr8E,90
-onnx_diagnostic/torch_models/validate.py,sha256=XNGZi7qSSytUczDfJ-X2ff5xvGdWdWkwjyz8ejxUSCE,79107
+onnx_diagnostic/torch_models/validate.py,sha256=0KL1vQmB9DTFuJqc8_CyddIztuwFx9qpRRPULHd-C04,80434
 onnx_diagnostic/torch_models/hghub/__init__.py,sha256=vi1Q7YHdddj1soiBN42MSvJdFqe2_KUoWafHISjwOu8,58
 onnx_diagnostic/torch_models/hghub/hub_api.py,sha256=rFbiPNLET-KdBpnv-p0nKgwHX6d7C_Z0s9zZ86_92kQ,14307
 onnx_diagnostic/torch_models/hghub/hub_data.py,sha256=8V_pAgACPLPsLRYUododg7MSL6str-T3tBEGY4OaeYQ,8724
 onnx_diagnostic/torch_models/hghub/hub_data_cached_configs.py,sha256=aSa_7Rjider6ruqQ2-fXQMyyDS8VhB1xKxcPNk8qUeU,288776
-onnx_diagnostic/torch_models/hghub/model_inputs.py,sha256=FaNFmWmzAqQQ7nM-L0eypeHG-YmCReXxwOOAb3UN7D0,15493
+onnx_diagnostic/torch_models/hghub/model_inputs.py,sha256=xIY_CWOp3m5-cJUvDLTZiH9GwiXi6xTYwONgFY4o45g,15593
 onnx_diagnostic/torch_models/hghub/model_specific.py,sha256=j50Nu7wddJMoqmD4QzMbNdFDUUgUmSBKRzPDH55TlUQ,2498
 onnx_diagnostic/torch_models/untrained/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 onnx_diagnostic/torch_models/untrained/llm_phi2.py,sha256=JbGZmW41MPJcQgqaJc9R2G00nI79nI-lABN-ffA1lmY,4037
@@ -121,8 +121,8 @@ onnx_diagnostic/torch_models/untrained/llm_tiny_llm.py,sha256=QXw_Bs2SzfeiQMf-tm
 onnx_diagnostic/torch_onnx/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 onnx_diagnostic/torch_onnx/runtime_info.py,sha256=1g9F_Jf9AAgYQU4stbsrFXwQl-30mWlQrFbQ7val8Ps,9268
 onnx_diagnostic/torch_onnx/sbs.py,sha256=IoKLA5UwS6kY8g4OOf_bdQwCziIsQfBczZ3w8wo4wZM,16905
-onnx_diagnostic-0.7.14.dist-info/licenses/LICENSE.txt,sha256=Vv6TXglX6Rc0d-f8aREhayhT-6PMQXEyOmI2NKlUCMc,1045
-onnx_diagnostic-0.7.14.dist-info/METADATA,sha256=id7f09epUAspAc4BxIlxRp0HhfGpR4SXI3BnYx0bjts,6730
-onnx_diagnostic-0.7.14.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-onnx_diagnostic-0.7.14.dist-info/top_level.txt,sha256=KwNkXewmcobM3ZT1DJLVWH6ebJzA5qKg7cWqKfpGNT4,16
-onnx_diagnostic-0.7.14.dist-info/RECORD,,
+onnx_diagnostic-0.7.15.dist-info/licenses/LICENSE.txt,sha256=Vv6TXglX6Rc0d-f8aREhayhT-6PMQXEyOmI2NKlUCMc,1045
+onnx_diagnostic-0.7.15.dist-info/METADATA,sha256=8PCb8jeG1AwC10iaBQRqNBE_JF7huNn2o-l_7BnwzzE,6730
+onnx_diagnostic-0.7.15.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+onnx_diagnostic-0.7.15.dist-info/top_level.txt,sha256=KwNkXewmcobM3ZT1DJLVWH6ebJzA5qKg7cWqKfpGNT4,16
+onnx_diagnostic-0.7.15.dist-info/RECORD,,