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

onnx_diagnostic/torch_export_patches/patches/_patch_transformers_funnel.py

@@ -0,0 +1,80 @@
+import torch
+
+try:
+    import transformers.models.funnel.modeling_funnel
+
+    patch_funnel = True
+except ImportError:
+    patch_funnel = False
+
+if patch_funnel:
+    from transformers.models.funnel.modeling_funnel import _relative_shift_gather
+
+    class patched_FunnelAttentionStructure(torch.nn.Module):
+        _PATCHES_ = ["relative_pos"]
+        _PATCHED_CLASS_ = transformers.models.funnel.modeling_funnel.FunnelAttentionStructure
+
+        def relative_pos(
+            self, pos: torch.Tensor, stride: int, pooled_pos=None, shift: int = 1
+        ) -> torch.Tensor:
+            if pooled_pos is None:
+                pooled_pos = pos
+            ref_point = pooled_pos[0] - pos[0]
+            # PATCHED
+            num_remove = shift * pooled_pos.shape[0]
+            max_dist = ref_point + num_remove * stride
+            min_dist = pooled_pos[0] - pos[-1]
+            return torch.arange(
+                max_dist.to(torch.long),
+                (min_dist - 1).to(torch.long),
+                torch.tensor(-stride, dtype=torch.long),
+                dtype=torch.long,
+                device=pos.device,
+            )
+
+    class patched_FunnelRelMultiheadAttention(torch.nn.Module):
+        _PATCHES_ = ["relative_positional_attention"]
+        _PATCHED_CLASS_ = (
+            transformers.models.funnel.modeling_funnel.FunnelRelMultiheadAttention
+        )
+
+        def relative_positional_attention(
+            self, position_embeds, q_head, context_len, cls_mask=None
+        ):
+            """Relative attention score for the positional encodings"""
+            # q_head has shape batch_size x sea_len x n_head x d_head
+            if self.config.attention_type == "factorized":
+                phi, pi, psi, omega = position_embeds
+                # Shape n_head x d_head
+                u = self.r_r_bias * self.scale
+                # Shape d_model x n_head x d_head
+                w_r = self.r_kernel
+
+                # Shape batch_size x sea_len x n_head x d_model
+                q_r_attention = torch.einsum("binh,dnh->bind", q_head + u, w_r)
+                q_r_attention_1 = q_r_attention * phi[:, None]
+                q_r_attention_2 = q_r_attention * pi[:, None]
+
+                # Shape batch_size x n_head x seq_len x context_len
+                positional_attn = torch.einsum(
+                    "bind,jd->bnij", q_r_attention_1, psi
+                ) + torch.einsum("bind,jd->bnij", q_r_attention_2, omega)
+            else:
+                shift = 2 if q_head.shape[1] != context_len else 1
+                r = position_embeds[self.block_index][shift - 1]
+                # Shape n_head x d_head
+                v = self.r_r_bias * self.scale
+                # Shape d_model x n_head x d_head
+                w_r = self.r_kernel
+
+                # Shape max_rel_len x n_head x d_model
+                r_head = torch.einsum("td,dnh->tnh", r, w_r)
+                # Shape batch_size x n_head x seq_len x max_rel_len
+                positional_attn = torch.einsum("binh,tnh->bnit", q_head + v, r_head)
+                # Shape batch_size x n_head x seq_len x context_len
+                positional_attn = _relative_shift_gather(positional_attn, context_len, shift)
+
+            if cls_mask is not None:
+                # PATCHED
+                positional_attn = positional_attn * cls_mask
+            return positional_attn

onnx_diagnostic/torch_export_patches/patches/_patch_transformers_qwen2_5.py

@@ -256,8 +256,23 @@ if patch_qwen2_5:
         return attn_output
 
     def qwen_version_selector(opset: int, *args: torch.Tensor) -> Tuple[str, torch.dtype]:
-        first_tensor = next(a for a in args if a is not None)
-        dtype = first_tensor.dtype
+        import onnx_ir
+
+        first_float_tensor = next(
+            a
+            for a in args
+            if a is not None
+            and a.dtype
+            in {
+                torch.float16,
+                torch.float32,
+                torch.bfloat16,
+                onnx_ir.DataType.BFLOAT16,
+                onnx_ir.DataType.FLOAT16,
+                onnx_ir.DataType.FLOAT,
+            }
+        )
+        dtype = first_float_tensor.dtype
         strategy = patched_Qwen2_5_VLVisionAttention.STRATEGY_FOR_ATTENTION()
         itype = torch_dtype_to_onnx_dtype(dtype)
         if strategy is not None:
@@ -269,7 +284,7 @@ if patch_qwen2_5:
         if dtype == torch.float16 or itype == onnx.TensorProto.FLOAT16:
             # first_tensor may be a SymbolicTensor (onnx).
             # is_cuda is not available.
-            if hasattr(first_tensor, "is_cuda") and first_tensor.is_cuda:
+            if hasattr(first_float_tensor, "is_cuda") and first_float_tensor.is_cuda:
                 return "PACKED", itype
             return "LOOPMHA", itype
         raise AssertionError(
@@ -733,3 +748,71 @@ if patch_qwen2_5:
             attn_output = attn_output.reshape(seq_length, -1).contiguous()
             attn_output = self.proj(attn_output)
             return attn_output
+
+    class patched_Qwen2_5_VLModel:
+        _PATCHES_ = ["get_placeholder_mask"]
+        _PATCHED_CLASS_ = transformers.models.qwen2_5_vl.modeling_qwen2_5_vl.Qwen2_5_VLModel
+
+        def get_placeholder_mask(
+            self,
+            input_ids: torch.LongTensor,
+            inputs_embeds: torch.FloatTensor,
+            image_features: Optional[torch.FloatTensor] = None,
+            video_features: Optional[torch.FloatTensor] = None,
+        ):
+            if input_ids is None:
+                special_image_mask = inputs_embeds == self.get_input_embeddings()(
+                    torch.tensor(
+                        self.config.image_token_id,
+                        dtype=torch.long,
+                        device=inputs_embeds.device,
+                    )
+                )
+                special_image_mask = special_image_mask.all(-1)
+                special_video_mask = inputs_embeds == self.get_input_embeddings()(
+                    torch.tensor(
+                        self.config.video_token_id,
+                        dtype=torch.long,
+                        device=inputs_embeds.device,
+                    )
+                )
+                special_video_mask = special_video_mask.all(-1)
+            else:
+                special_image_mask = input_ids == self.config.image_token_id
+                special_video_mask = input_ids == self.config.video_token_id
+
+            special_image_mask = (
+                special_image_mask.unsqueeze(-1)
+                .expand_as(inputs_embeds)
+                .to(inputs_embeds.device)
+            )
+
+            # PATCHED: we should use torch._check
+            # but this fails for compilation. It cannot be verified with FakeTensors
+            # torch._check(
+            #    image_features is None
+            #    or inputs_embeds[special_image_mask].numel() == image_features.numel(),
+            #    lambda: (
+            #        f"Image features and image tokens do not match: tokens: "
+            #        f"{special_image_mask.sum()}, features {image_features.shape[0]}"
+            #    ),
+            # )
+
+            special_video_mask = (
+                special_video_mask.unsqueeze(-1)
+                .expand_as(inputs_embeds)
+                .to(inputs_embeds.device)
+            )
+
+            # PATCHED: we should use torch._check
+            # but this fails for compilation. It cannot be verified with FakeTensors
+            # torch._check(
+            #    video_features is None
+            #    or inputs_embeds[special_video_mask].numel() == video_features.numel(),
+            #    lambda: (
+            #        f"Videos features and video tokens do not match: tokens: "
+            #        f"{special_video_mask.sum()}, features {video_features.shape[0]}"
+            #    ),
+            # )
+
+            return special_image_mask, special_video_mask

onnx_diagnostic/torch_export_patches/patches/patch_torch.py

@@ -5,6 +5,7 @@ import os
 import traceback
 from functools import reduce
 from typing import Any, Callable, cast, Dict, List, Optional, Sequence, Tuple, Union
+import sympy
 import torch
 from torch._subclasses.fake_tensor import FakeTensorMode
 
@@ -1091,3 +1092,17 @@ def patched__broadcast_in_dim_meta_level_2(
                 new_strides.append(a.stride()[original_idx] * a.size()[original_idx])
 
     return a.as_strided(shape, new_strides, a.storage_offset())
+
+
+class patched_DynamicDimConstraintPrinter:
+    """
+    Patches
+    ``torch.tx.experimental.symbolic_shapes.DynamicDimConstraintPrinter._print_Symbol``.
+    Valid for ``torch>=2.10``.
+    """
+
+    def _print_Symbol(self, expr: sympy.Symbol) -> str:
+        assert isinstance(expr, sympy.Symbol), str(type(expr))
+        if self.symbol_to_source.get(expr):
+            return self.symbol_to_source[expr][0].name
+        return str(expr)

onnx_diagnostic/torch_export_patches/patches/patch_transformers.py

@@ -1,29 +1,37 @@
 # transformers
 from typing import List
 from .patch_helper import _has_transformers
-
 from ._patch_transformers_attention import (
     patched_sdpa_attention_forward,
     patched_model_bart_eager_attention_forward,
     patched_modeling_marian_eager_attention_forward,
 )
+from ._patch_transformers_generation_mixin import patched_GenerationMixin
+from ._patch_transformers_causal_mask import patched_AttentionMaskConverter
+from ._patch_transformers_rotary_embedding import (
+    patched__compute_dynamic_ntk_parameters,
+    patched_dynamic_rope_update,
+    patched_GemmaRotaryEmbedding,
+    patched_LlamaRotaryEmbedding,
+    patched_MistralRotaryEmbedding,
+    patched_MixtralRotaryEmbedding,
+    patched_PhiRotaryEmbedding,
+)
+from ._patch_transformers_idefics import patched_IdeficsEmbedding, patched_IdeficsAttention
+from ._patch_transformers_sam_mask_decoder import patched_SamMaskDecoder
+
+# transformers dependent patches
 
 from ._patch_transformers_cache_utils import patch_parse_processor_args
 
 if patch_parse_processor_args:
     from ._patch_transformers_cache_utils import patched_parse_processor_args
-
-from ._patch_transformers_causal_mask import patched_AttentionMaskConverter
-
 from ._patch_transformers_dynamic_cache import patch_DynamicLayer, patch_DynamicCache
 
 if patch_DynamicLayer:
     from ._patch_transformers_dynamic_cache import patched_DynamicLayer
 if patch_DynamicCache:
     from ._patch_transformers_dynamic_cache import patched_DynamicCache
-
-from ._patch_transformers_generation_mixin import patched_GenerationMixin
-
 from ._patch_transformers_masking_utils import patch_masking_utils
 
 if patch_masking_utils:
@@ -33,15 +41,7 @@ if patch_masking_utils:
         patched_sdpa_mask_recent_torch,
     )
 
-from ._patch_transformers_rotary_embedding import (
-    patched__compute_dynamic_ntk_parameters,
-    patched_dynamic_rope_update,
-    patched_GemmaRotaryEmbedding,
-    patched_LlamaRotaryEmbedding,
-    patched_MistralRotaryEmbedding,
-    patched_MixtralRotaryEmbedding,
-    patched_PhiRotaryEmbedding,
-)
+# transformers models dependent patches
 
 if _has_transformers("4.51"):
     from ._patch_transformers_rotary_embedding import patched_Phi3RotaryEmbedding
@@ -54,16 +54,11 @@ if _has_transformers("4.52"):
 if _has_transformers("4.53"):
     from ._patch_transformers_rotary_embedding import patched_SmolLM3RotaryEmbedding
 
-# Models
-
 from ._patch_transformers_gemma3 import patch_gemma3
 
 if patch_gemma3:
     from ._patch_transformers_gemma3 import patched_Gemma3Model
 
-from ._patch_transformers_idefics import patched_IdeficsEmbedding, patched_IdeficsAttention
-
-
 from ._patch_transformers_qwen2 import patch_qwen2
 
 if patch_qwen2:
@@ -77,16 +72,20 @@ if patch_qwen2_5:
         patched_Qwen2_5_VisionTransformerPretrainedModel,
         patched_Qwen2_5_VLVisionAttentionOneIteration,
         patched_Qwen2_5_VLVisionAttention,
+        patched_Qwen2_5_VLModel,
         PLUGS as PLUGS_Qwen25,
     )
-
 from ._patch_transformers_qwen3 import patch_qwen3
 
 if patch_qwen3:
     from ._patch_transformers_qwen3 import patched_Qwen3MoeSparseMoeBlock
+from ._patch_transformers_funnel import patch_funnel
 
-
-from ._patch_transformers_sam_mask_decoder import patched_SamMaskDecoder
+if patch_funnel:
+    from ._patch_transformers_funnel import (
+        patched_FunnelAttentionStructure,
+        patched_FunnelRelMultiheadAttention,
+    )
 
 
 def get_transformers_plugs() -> List["EagerDirectReplacementWithOnnx"]:  # noqa: F821

onnx_diagnostic/torch_models/hghub/hub_api.py

@@ -184,7 +184,18 @@ def _trygetattr(config, attname):
         return None
 
 
+def rewrite_architecture_name(name: Optional[str]) -> Optional[str]:
+    if name == "ConditionalDETRForObjectDetection":
+        return "ConditionalDetrForObjectDetection"
+    return name
+
+
 def architecture_from_config(config) -> Optional[str]:
+    """Guesses the architecture (class) of the model described by this config."""
+    return rewrite_architecture_name(_architecture_from_config(config))
+
+
+def _architecture_from_config(config) -> Optional[str]:
     """Guesses the architecture (class) of the model described by this config."""
     if isinstance(config, dict):
         if "_class_name" in config:

onnx_diagnostic/torch_models/hghub/hub_data.py

@@ -5,7 +5,10 @@ from typing import Dict, List
 
 __date__ = "2025-06-21"
 
-__data_arch_values__ = {"ResNetForImageClassification": dict(image_size=224)}
+__data_arch_values__ = {
+    "ConditionalDETRForObjectDetection": dict(image_size=224),
+    "ResNetForImageClassification": dict(image_size=224),
+}
 
 __data_arch__ = textwrap.dedent(
     """
@@ -32,6 +35,7 @@ __data_arch__ = textwrap.dedent(
     ConvNextV2Model,image-feature-extraction
     CosmosTransformer3DModel,image-to-video
     CvtModel,feature-extraction
+    ClvpModelForConditionalGeneration,audio-feature-extraction
     DPTModel,image-feature-extraction
     Data2VecAudioModel,feature-extraction
     Data2VecTextModel,feature-extraction
@@ -49,6 +53,8 @@ __data_arch__ = textwrap.dedent(
     ElectraModel,feature-extraction
     EsmModel,feature-extraction
     FalconMambaForCausalLM,text-generation
+    FunnelBaseModel,feature-extraction
+    FuyuForCausalLM,image-text-to-text
     GLPNModel,image-feature-extraction
     GPT2LMHeadModel,text-generation
     GPTBigCodeModel,feature-extraction
@@ -63,6 +69,7 @@ __data_arch__ = textwrap.dedent(
     Glm4vMoeForConditionalGeneration,image-text-to-text
     GraniteForCausalLM,text-generation
     GroupViTModel,feature-extraction
+    HeliumForCausalLM,text-generation
     HieraForImageClassification,image-classification
     HubertModel,feature-extraction
     IBertModel,feature-extraction
@@ -136,6 +143,7 @@ __data_arch__ = textwrap.dedent(
     SwinModel,image-feature-extraction
     Swinv2Model,image-feature-extraction
     T5ForConditionalGeneration,text2text-generation
+    T5GemmaForConditionalGeneration,text2text-generation
     TableTransformerModel,image-feature-extraction
     TableTransformerForObjectDetection,object-detection
     UNet2DConditionModel,text-to-image

onnx_diagnostic/torch_models/hghub/hub_data_cached_configs.py

@@ -55,6 +55,7 @@ Automatically generated:
 import base64
 import json
 import textwrap
+from typing import Any
 import transformers
 
 null = None
@@ -62,6 +63,22 @@ true = True
 false = False
 
 
+def _enforce_default(config_type: type, **kwargs) -> Any:
+    config = config_type(**kwargs)
+    for name in [
+        *[k for k in kwargs if k.endswith("_token_id")],
+        "attention_dropout",
+        "hidden_size",
+        "hidden_act",
+        "intermediate_size",
+        "max_position_embeddings",
+        "vocab_size",
+    ]:
+        if name in kwargs and (not hasattr(config, name) or getattr(config, name) is None):
+            setattr(config, name, kwargs[name])
+    return config
+
+
 def _ccached_arnir0_tiny_LLM():
     "arnir0/Tiny-LLM"
     return transformers.LlamaConfig(
@@ -4691,7 +4708,8 @@ def _ccached_zai_glm_45():
 
 def _ccached_microsoft_phi3_mini_128k_instruct():
     "microsoft/Phi-3-mini-128k-instruct"
-    return transformers.Phi3Config(
+    return _enforce_default(
+        transformers.Phi3Config,
         **{
             "_name_or_path": "Phi-3-mini-128k-instruct",
             "architectures": ["Phi3ForCausalLM"],
@@ -4827,13 +4845,14 @@ def _ccached_microsoft_phi3_mini_128k_instruct():
             "use_cache": true,
             "attention_bias": false,
             "vocab_size": 32064,
-        }
+        },
     )
 
 
 def _ccached_google_gemma_3_4b_it_like():
     "google/gemma-3-4b-it"
-    return transformers.Gemma3Config(
+    return _enforce_default(
+        transformers.Gemma3Config,
         **{
             "architectures": ["Gemma3ForConditionalGeneration"],
             "boi_token_index": 255999,
@@ -4863,13 +4882,14 @@ def _ccached_google_gemma_3_4b_it_like():
                 "patch_size": 14,
                 "vision_use_head": false,
             },
-        }
+        },
     )
 
 
 def _ccached_hf_internal_testing_tiny_random_gemma3_for_causal_lm():
     "hf-internal-testing/tiny-random-Gemma3ForCausalLM"
-    return transformers.Gemma3TextConfig(
+    return _enforce_default(
+        transformers.Gemma3TextConfig,
         **{
             "architectures": ["Gemma3ForCausalLM"],
             "attention_bias": false,
@@ -4901,13 +4921,14 @@ def _ccached_hf_internal_testing_tiny_random_gemma3_for_causal_lm():
             "transformers_version": "4.52.0.dev0",
             "use_cache": true,
             "vocab_size": 262144,
-        }
+        },
     )
 
 
 def _ccached_qwen_qwen2_5_vl_7b_instruct():
     "Qwen/Qwen2.5-VL-7B-Instruct"
-    return transformers.Qwen2_5_VLConfig(
+    return _enforce_default(
+        transformers.Qwen2_5_VLConfig,
         **{
             "architectures": ["Qwen2_5_VLForConditionalGeneration"],
             "attention_dropout": 0.0,
@@ -4954,5 +4975,5 @@ def _ccached_qwen_qwen2_5_vl_7b_instruct():
             },
             "rope_scaling": {"type": "mrope", "mrope_section": [16, 24, 24]},
             "vocab_size": 152064,
-        }
+        },
     )

onnx_diagnostic/torch_models/hghub/model_inputs.py

@@ -64,6 +64,7 @@ def get_untrained_model_with_inputs(
     use_only_preinstalled: bool = False,
     config_reduction: Optional[Callable[[Any, str], Dict]] = None,
     submodule: Optional[str] = None,
+    skip_inputs: bool = False,
 ) -> Dict[str, Any]:
     """
     Gets a non initialized model similar to the original model
@@ -93,6 +94,7 @@ def get_untrained_model_with_inputs(
         this function takes a configuration and a task (string)
         as arguments
     :param submodule: use a submodule instead of the main model
+    :param skip_inputs: do not generate the inputs
     :return: dictionary with a model, inputs, dynamic shapes, and the configuration,
         some necessary rewriting as well
 
@@ -332,13 +334,12 @@ def get_untrained_model_with_inputs(
                         f"[get_untrained_model_with_inputs] "
                         f"instantiate_specific_model(2) {cls_model}"
                     )
-
                 try:
                     if type(config) is dict:
                         model = cls_model(**config)
                     else:
                         model = cls_model(config)
-                except RuntimeError as e:
+                except (RuntimeError, AttributeError, ValueError) as e:
                     raise RuntimeError(
                         f"Unable to instantiate class {cls_model.__name__} with\n{config}"
                     ) from e
@@ -350,23 +351,27 @@ def get_untrained_model_with_inputs(
                     )
 
     # input kwargs
-    seed = int(os.environ.get("SEED", "17")) + 1
-    torch.manual_seed(seed)
-    kwargs, fct = random_input_kwargs(config, task)  # type: ignore[arg-type]
-    if verbose:
-        print(f"[get_untrained_model_with_inputs] use fct={fct}")
-        if os.environ.get("PRINT_CONFIG") in (1, "1"):
-            print(f"-- input kwargs for task {task!r}")
-            pprint.pprint(kwargs)
-    if inputs_kwargs:
-        kwargs.update(inputs_kwargs)
-
-    # This line is important. Some models may produce different
-    # outputs even with the same inputs in training mode.
-    model.eval()  # type: ignore[union-attr]
-    res = fct(model, config, add_second_input=add_second_input, **kwargs)
-
-    res["input_kwargs"] = kwargs
+    if not skip_inputs:
+        seed = int(os.environ.get("SEED", "17")) + 1
+        torch.manual_seed(seed)
+        kwargs, fct = random_input_kwargs(config, task)  # type: ignore[arg-type]
+        if verbose:
+            print(f"[get_untrained_model_with_inputs] use fct={fct}")
+            if os.environ.get("PRINT_CONFIG") in (1, "1"):
+                print(f"-- input kwargs for task {task!r}")
+                pprint.pprint(kwargs)
+        if inputs_kwargs:
+            kwargs.update(inputs_kwargs)
+
+        # This line is important. Some models may produce different
+        # outputs even with the same inputs in training mode.
+        model.eval()  # type: ignore[union-attr]
+        res = fct(model, config, add_second_input=add_second_input, **kwargs)
+
+        res["input_kwargs"] = kwargs
+    else:
+        res = {}
+
     res["model_kwargs"] = mkwargs
     if diff_config is not None:
         res["dump_info"] = dict(config_diff=diff_config)