onnx-diagnostic 0.7.5__py3-none-any.whl → 0.7.7__py3-none-any.whl

onnx_diagnostic/torch_export_patches/serialization/transformers_impl.py

@@ -1,15 +1,25 @@
 from typing import Any, List, Set, Tuple
 import torch
-import transformers
 from transformers.cache_utils import (
     DynamicCache,
-    MambaCache,
     EncoderDecoderCache,
+    HybridCache,
     SlidingWindowCache,
     StaticCache,
 )
+
+try:
+    from transformers.models.mamba.modeling_mamba import MambaCache
+except ImportError:
+    from transformers.cache_utils import MambaCache
 from transformers.modeling_outputs import BaseModelOutput
-from ...helpers.cache_helper import make_static_cache
+from ...helpers.cache_helper import (
+    make_dynamic_cache,
+    make_hybrid_cache,
+    make_sliding_window_cache,
+    make_static_cache,
+    CacheKeyValue,
+)
 from . import make_serialization_function_for_dataclass
 
 
@@ -29,6 +39,12 @@ def flatten_mamba_cache(
     mamba_cache: MambaCache,
 ) -> Tuple[List[Any], torch.utils._pytree.Context]:
     """Serializes a :class:`transformers.cache_utils.MambaCache` with python objects."""
+    assert isinstance(mamba_cache.conv_states, list) and isinstance(
+        mamba_cache.ssm_states, list
+    ), (
+        f"Unexpected types for conv_states and ssm_states {type(mamba_cache.conv_states)}, "
+        f"{type(mamba_cache.ssm_states)}"
+    )
     flat = [
         ("conv_states", mamba_cache.conv_states),
         ("ssm_states", mamba_cache.ssm_states),
@@ -85,9 +101,8 @@ def flatten_dynamic_cache(
     dynamic_cache: DynamicCache,
 ) -> Tuple[List[Any], torch.utils._pytree.Context]:
     """Serializes a :class:`transformers.cache_utils.DynamicCache` with python objects."""
-    if hasattr(transformers.cache_utils, "_flatten_dynamic_cache"):
-        return transformers.cache_utils._flatten_dynamic_cache(dynamic_cache)
-    flat = [("key_cache", dynamic_cache.key_cache), ("value_cache", dynamic_cache.value_cache)]
+    ca = CacheKeyValue(dynamic_cache)
+    flat = [("key_cache", ca.key_cache), ("value_cache", ca.value_cache)]
     return [f[1] for f in flat], [f[0] for f in flat]
 
 
@@ -95,8 +110,6 @@ def flatten_with_keys_dynamic_cache(
     dynamic_cache: DynamicCache,
 ) -> Tuple[List[Tuple[torch.utils._pytree.KeyEntry, Any]], torch.utils._pytree.Context]:
     """Serializes a :class:`transformers.cache_utils.DynamicCache` with python objects."""
-    if hasattr(transformers.cache_utils, "_flatten_with_keys_dynamic_cache"):
-        return transformers.cache_utils._flatten_with_keys_dynamic_cache(dynamic_cache)
     values, context = flatten_dynamic_cache(dynamic_cache)
     return [(torch.utils._pytree.MappingKey(k), v) for k, v in zip(context, values)], context
 
@@ -105,15 +118,36 @@ def unflatten_dynamic_cache(
     values: List[Any], context: torch.utils._pytree.Context, output_type=None
 ) -> DynamicCache:
     """Restores a :class:`transformers.cache_utils.DynamicCache` from python objects."""
-    if hasattr(transformers.cache_utils, "_unflatten_dynamic_cache"):
-        assert output_type is None, f"output_type={output_type} not supported"
-        return transformers.cache_utils._unflatten_dynamic_cache(values, context)
+    return make_dynamic_cache(list(zip(values[0], values[1])))
 
-    cache = transformers.cache_utils.DynamicCache()
-    values = dict(zip(context, values))
-    for k, v in values.items():
-        setattr(cache, k, v)
-    return cache
+
+#############
+# HybridCache
+#############
+
+
+def flatten_hybrid_cache(
+    cache: HybridCache,
+) -> Tuple[List[Any], torch.utils._pytree.Context]:
+    """Serializes a :class:`transformers.cache_utils.HybridCache` with python objects."""
+    ca = CacheKeyValue(cache)
+    flat = [("key_cache", ca.key_cache), ("value_cache", ca.value_cache)]
+    return [f[1] for f in flat], [f[0] for f in flat]
+
+
+def flatten_with_keys_hybrid_cache(
+    cache: HybridCache,
+) -> Tuple[List[Tuple[torch.utils._pytree.KeyEntry, Any]], torch.utils._pytree.Context]:
+    """Serializes a :class:`transformers.cache_utils.HybridCache` with python objects."""
+    values, context = flatten_hybrid_cache(cache)
+    return [(torch.utils._pytree.MappingKey(k), v) for k, v in zip(context, values)], context
+
+
+def unflatten_hybrid_cache(
+    values: List[Any], context: torch.utils._pytree.Context, output_type=None
+) -> HybridCache:
+    """Restores a :class:`transformers.cache_utils.HybridCache` from python objects."""
+    return make_hybrid_cache(list(zip(values[0], values[1])))
 
 
 #############
@@ -125,12 +159,13 @@ def flatten_static_cache(
     cache: StaticCache,
 ) -> Tuple[List[Any], torch.utils._pytree.Context]:
     """Serializes a :class:`transformers.cache_utils.StaticCache` with python objects."""
-    assert not cache.key_cache or cache.max_cache_len == cache.key_cache[0].shape[2], (
+    ca = CacheKeyValue(cache)
+    assert not ca.key_cache or cache.max_cache_len == ca.key_cache[0].shape[2], (
         f"Serialization doet not work when "
         f"cache.max_cache_len={cache.max_cache_len} != "
-        f"cache.key_cache[0].shape[2]={cache.key_cache[0].shape[2]}"
+        f"cache.key_cache[0].shape[2]={ca.keu_cache[0].shape[2]}"
     )
-    flat = [("key_cache", cache.key_cache), ("value_cache", cache.value_cache)]
+    flat = [("key_cache", ca.key_cache), ("value_cache", ca.value_cache)]
     return [f[1] for f in flat], [f[0] for f in flat]
 
 
@@ -163,7 +198,8 @@ def flatten_sliding_window_cache(
     Serializes a :class:`transformers.cache_utils.SlidingWindowCache`
     with python objects.
     """
-    flat = [("key_cache", cache.key_cache), ("value_cache", cache.value_cache)]
+    ca = CacheKeyValue(cache)
+    flat = [("key_cache", ca.key_cache), ("value_cache", ca.value_cache)]
     return [f[1] for f in flat], [f[0] for f in flat]
 
 
@@ -183,26 +219,7 @@ def unflatten_sliding_window_cache(
 ) -> SlidingWindowCache:
     """Restores a :class:`transformers.cache_utils.SlidingWindowCache` from python objects."""
     key_cache, value_cache = values
-
-    class _config:
-        def __init__(self):
-            self.head_dim = key_cache[0].shape[-1]
-            self.num_attention_heads = key_cache[0].shape[1]
-            self.num_hidden_layers = len(key_cache)
-            self.sliding_window = key_cache[0].shape[2]
-
-    cache = SlidingWindowCache(
-        _config(),
-        max_batch_size=key_cache[0].shape[0],
-        max_cache_len=key_cache[0].shape[2],  # sligding window
-        device=key_cache[0].device,
-        dtype=key_cache[0].dtype,
-    )
-
-    values = dict(zip(context, values))
-    for k, v in values.items():
-        setattr(cache, k, v)
-    return cache
+    return make_sliding_window_cache(list(zip(values[0], values[1])))
 
 
 #####################
@@ -244,7 +261,9 @@ def unflatten_encoder_decoder_cache(
 ) -> EncoderDecoderCache:
     """Restores a :class:`transformers.cache_utils.EncoderDecoderCache` from python objects."""
     dictionary = torch.utils._pytree._dict_unflatten(values, context)
-    return EncoderDecoderCache(**dictionary)
+    return EncoderDecoderCache(
+        dictionary["self_attention_cache"], dictionary["cross_attention_cache"]
+    )
 
 
 #############

onnx_diagnostic/torch_models/hghub/hub_data.py

@@ -37,6 +37,7 @@ __data_arch__ = textwrap.dedent(
     DebertaModel,feature-extraction
     DebertaV2Model,feature-extraction
     DecisionTransformerModel,reinforcement-learning
+    DeepseekV3ForCausalLM,text-generation
     DeiTModel,image-feature-extraction
     DetrModel,image-feature-extraction
     Dinov2Model,image-feature-extraction
@@ -52,9 +53,12 @@ __data_arch__ = textwrap.dedent(
     GPTJModel,feature-extraction
     GPTNeoModel,feature-extraction
     GPTNeoXForCausalLM,text-generation
+    GptOssForCausalLM,text-generation
     GemmaForCausalLM,text-generation
     Gemma2ForCausalLM,text-generation
     Gemma3ForConditionalGeneration,image-text-to-text
+    Gemma3ForCausalLM,text-generation
+    Glm4vMoeForConditionalGeneration,image-text-to-text
     GraniteForCausalLM,text-generation
     GroupViTModel,feature-extraction
     HieraForImageClassification,image-classification
@@ -107,6 +111,7 @@ __data_arch__ = textwrap.dedent(
     PvtForImageClassification,image-classification
     Qwen2ForCausalLM,text-generation
     Qwen2_5_VLForConditionalGeneration,image-text-to-text
+    Qwen3MoeForCausalLM,text-generation
     RTDetrForObjectDetection,object-detection
     RegNetModel,image-feature-extraction
     RemBertModel,feature-extraction

onnx_diagnostic/torch_models/hghub/hub_data_cached_configs.py

@@ -1366,6 +1366,236 @@ def _ccached_fxmarty_tiny_random_gemmaforcausallm():
     )
 
 
+def _ccached_fxmarty_sam_vit_tiny_random():
+    "fxmarty/sam-vit-tiny-random"
+    return transformers.SamConfig(
+        **{
+            "_commit_hash": "a7c34ea5d2b33a3bc34d34dc9a7b2417c0eaa809",
+            "_name_or_path": "facebook/sam-vit-base",
+            "architectures": ["SamModel"],
+            "initializer_range": 0.02,
+            "mask_decoder_config": {
+                "_name_or_path": "",
+                "add_cross_attention": false,
+                "architectures": null,
+                "attention_downsample_rate": 2,
+                "bad_words_ids": null,
+                "begin_suppress_tokens": null,
+                "bos_token_id": null,
+                "chunk_size_feed_forward": 0,
+                "cross_attention_hidden_size": null,
+                "decoder_start_token_id": null,
+                "diversity_penalty": 0.0,
+                "do_sample": false,
+                "early_stopping": false,
+                "encoder_no_repeat_ngram_size": 0,
+                "eos_token_id": null,
+                "exponential_decay_length_penalty": null,
+                "finetuning_task": null,
+                "forced_bos_token_id": null,
+                "forced_eos_token_id": null,
+                "hidden_act": "relu",
+                "hidden_size": 32,
+                "id2label": {"0": "LABEL_0", "1": "LABEL_1"},
+                "iou_head_depth": 3,
+                "iou_head_hidden_dim": 256,
+                "is_decoder": false,
+                "is_encoder_decoder": false,
+                "label2id": {"LABEL_0": 0, "LABEL_1": 1},
+                "layer_norm_eps": 1e-06,
+                "length_penalty": 1.0,
+                "max_length": 20,
+                "min_length": 0,
+                "mlp_dim": 2048,
+                "model_type": "",
+                "no_repeat_ngram_size": 0,
+                "num_attention_heads": 8,
+                "num_beam_groups": 1,
+                "num_beams": 1,
+                "num_hidden_layers": 2,
+                "num_multimask_outputs": 3,
+                "num_return_sequences": 1,
+                "output_attentions": false,
+                "output_hidden_states": false,
+                "output_scores": false,
+                "pad_token_id": null,
+                "prefix": null,
+                "problem_type": null,
+                "pruned_heads": {},
+                "remove_invalid_values": false,
+                "repetition_penalty": 1.0,
+                "return_dict": true,
+                "return_dict_in_generate": false,
+                "sep_token_id": null,
+                "suppress_tokens": null,
+                "task_specific_params": null,
+                "temperature": 1.0,
+                "tf_legacy_loss": false,
+                "tie_encoder_decoder": false,
+                "tie_word_embeddings": true,
+                "tokenizer_class": null,
+                "top_k": 50,
+                "top_p": 1.0,
+                "torch_dtype": null,
+                "torchscript": false,
+                "transformers_version": "4.29.0.dev0",
+                "typical_p": 1.0,
+                "use_bfloat16": false,
+            },
+            "model_type": "sam",
+            "prompt_encoder_config": {
+                "_name_or_path": "",
+                "add_cross_attention": false,
+                "architectures": null,
+                "bad_words_ids": null,
+                "begin_suppress_tokens": null,
+                "bos_token_id": null,
+                "chunk_size_feed_forward": 0,
+                "cross_attention_hidden_size": null,
+                "decoder_start_token_id": null,
+                "diversity_penalty": 0.0,
+                "do_sample": false,
+                "early_stopping": false,
+                "encoder_no_repeat_ngram_size": 0,
+                "eos_token_id": null,
+                "exponential_decay_length_penalty": null,
+                "finetuning_task": null,
+                "forced_bos_token_id": null,
+                "forced_eos_token_id": null,
+                "hidden_act": "gelu",
+                "hidden_size": 32,
+                "id2label": {"0": "LABEL_0", "1": "LABEL_1"},
+                "image_embedding_size": 64,
+                "image_size": 1024,
+                "is_decoder": false,
+                "is_encoder_decoder": false,
+                "label2id": {"LABEL_0": 0, "LABEL_1": 1},
+                "layer_norm_eps": 1e-06,
+                "length_penalty": 1.0,
+                "mask_input_channels": 16,
+                "max_length": 20,
+                "min_length": 0,
+                "model_type": "",
+                "no_repeat_ngram_size": 0,
+                "num_beam_groups": 1,
+                "num_beams": 1,
+                "num_point_embeddings": 4,
+                "num_return_sequences": 1,
+                "output_attentions": false,
+                "output_hidden_states": false,
+                "output_scores": false,
+                "pad_token_id": null,
+                "patch_size": 16,
+                "prefix": null,
+                "problem_type": null,
+                "pruned_heads": {},
+                "remove_invalid_values": false,
+                "repetition_penalty": 1.0,
+                "return_dict": true,
+                "return_dict_in_generate": false,
+                "sep_token_id": null,
+                "suppress_tokens": null,
+                "task_specific_params": null,
+                "temperature": 1.0,
+                "tf_legacy_loss": false,
+                "tie_encoder_decoder": false,
+                "tie_word_embeddings": true,
+                "tokenizer_class": null,
+                "top_k": 50,
+                "top_p": 1.0,
+                "torch_dtype": null,
+                "torchscript": false,
+                "transformers_version": "4.29.0.dev0",
+                "typical_p": 1.0,
+                "use_bfloat16": false,
+            },
+            "torch_dtype": "float32",
+            "transformers_version": null,
+            "vision_config": {
+                "_name_or_path": "",
+                "add_cross_attention": false,
+                "architectures": null,
+                "attention_dropout": 0.0,
+                "bad_words_ids": null,
+                "begin_suppress_tokens": null,
+                "bos_token_id": null,
+                "chunk_size_feed_forward": 0,
+                "cross_attention_hidden_size": null,
+                "decoder_start_token_id": null,
+                "diversity_penalty": 0.0,
+                "do_sample": false,
+                "dropout": 0.0,
+                "early_stopping": false,
+                "encoder_no_repeat_ngram_size": 0,
+                "eos_token_id": null,
+                "exponential_decay_length_penalty": null,
+                "finetuning_task": null,
+                "forced_bos_token_id": null,
+                "forced_eos_token_id": null,
+                "global_attn_indexes": [2, 5, 8, 11],
+                "hidden_act": "gelu",
+                "hidden_size": 96,
+                "id2label": {"0": "LABEL_0", "1": "LABEL_1"},
+                "image_size": 1024,
+                "initializer_factor": 1.0,
+                "initializer_range": 1e-10,
+                "intermediate_size": 768,
+                "is_decoder": false,
+                "is_encoder_decoder": false,
+                "label2id": {"LABEL_0": 0, "LABEL_1": 1},
+                "layer_norm_eps": 1e-06,
+                "length_penalty": 1.0,
+                "max_length": 20,
+                "min_length": 0,
+                "mlp_dim": 384,
+                "mlp_ratio": 4.0,
+                "model_type": "",
+                "no_repeat_ngram_size": 0,
+                "num_attention_heads": 1,
+                "num_beam_groups": 1,
+                "num_beams": 1,
+                "num_channels": 3,
+                "num_hidden_layers": 12,
+                "num_pos_feats": 16,
+                "num_return_sequences": 1,
+                "output_attentions": false,
+                "output_channels": 32,
+                "output_hidden_states": false,
+                "output_scores": false,
+                "pad_token_id": null,
+                "patch_size": 16,
+                "prefix": null,
+                "problem_type": null,
+                "projection_dim": 64,
+                "pruned_heads": {},
+                "qkv_bias": true,
+                "remove_invalid_values": false,
+                "repetition_penalty": 1.0,
+                "return_dict": true,
+                "return_dict_in_generate": false,
+                "sep_token_id": null,
+                "suppress_tokens": null,
+                "task_specific_params": null,
+                "temperature": 1.0,
+                "tf_legacy_loss": false,
+                "tie_encoder_decoder": false,
+                "tie_word_embeddings": true,
+                "tokenizer_class": null,
+                "top_k": 50,
+                "top_p": 1.0,
+                "torch_dtype": null,
+                "torchscript": false,
+                "transformers_version": "4.29.0.dev0",
+                "typical_p": 1.0,
+                "use_abs_pos": true,
+                "use_bfloat16": false,
+                "use_rel_pos": true,
+                "window_size": 14,
+            },
+        }
+    )
+
+
 def _ccached_hf_internal_testing_tiny_random_gptneoxforcausallm():
     "hf-internal-testing/tiny-random-GPTNeoXForCausalLM"
     return transformers.GPTNeoXConfig(
@@ -4330,3 +4560,61 @@ def _ccached_diffusers_tiny_torch_full_checker_unet():
         "up_block_types": ["CrossAttnUpBlock2D", "UpBlock2D"],
         "use_linear_projection": false,
     }
+
+
+def _ccached_riny_random_gemma_3():
+    "tiny-random/gemma-3"
+    return transformers.Gemma3Config(
+        **{
+            "architectures": ["Gemma3ForConditionalGeneration"],
+            "boi_token_index": 255999,
+            "eoi_token_index": 256000,
+            "eos_token_id": [1, 106],
+            "image_token_index": 262144,
+            "initializer_range": 0.02,
+            "mm_tokens_per_image": 256,
+            "model_type": "gemma3",
+            "text_config": {
+                "attention_bias": false,
+                "attention_dropout": 0.0,
+                "attn_logit_softcapping": null,
+                "cache_implementation": "hybrid",
+                "final_logit_softcapping": null,
+                "head_dim": 32,
+                "hidden_activation": "gelu_pytorch_tanh",
+                "hidden_size": 32,
+                "initializer_range": 0.02,
+                "intermediate_size": 128,
+                "max_position_embeddings": 131072,
+                "model_type": "gemma3_text",
+                "num_attention_heads": 1,
+                "num_hidden_layers": 2,
+                "num_key_value_heads": 1,
+                "query_pre_attn_scalar": 168,
+                "rms_norm_eps": 1e-06,
+                "rope_local_base_freq": 10000.0,
+                "rope_scaling": {"factor": 8.0, "rope_type": "linear"},
+                "rope_theta": 1000000.0,
+                "sliding_window": 1024,
+                "sliding_window_pattern": 2,
+                "use_cache": true,
+                "vocab_size": 262208,
+            },
+            "torch_dtype": "bfloat16",
+            "transformers_version": "4.50.0.dev0",
+            "vision_config": {
+                "attention_dropout": 0.0,
+                "hidden_act": "gelu_pytorch_tanh",
+                "hidden_size": 32,
+                "image_size": 896,
+                "intermediate_size": 128,
+                "layer_norm_eps": 1e-06,
+                "model_type": "siglip_vision_model",
+                "num_attention_heads": 1,
+                "num_channels": 3,
+                "num_hidden_layers": 2,
+                "patch_size": 14,
+                "vision_use_head": false,
+            },
+        }
+    )

onnx_diagnostic/torch_models/validate.py

@@ -426,6 +426,7 @@ 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}")
+        print(f"[validate_model] patch={patch!r}")
         if model_options:
             print(f"[validate_model] model_options={model_options!r}")
         print(f"[validate_model] get dummy inputs with input_options={input_options}...")

{onnx_diagnostic-0.7.5.dist-info → onnx_diagnostic-0.7.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: onnx-diagnostic
-Version: 0.7.5
+Version: 0.7.7
 Summary: Investigate ONNX models
 Home-page: https://github.com/sdpython/onnx-diagnostic
 Author: Xavier Dupré
@@ -27,7 +27,7 @@ Requires-Dist: numpy
 Requires-Dist: onnx>=1.16.0
 Requires-Dist: onnxruntime>=1.21
 Requires-Dist: optree
-Requires-Dist: torch>=2.7
+Requires-Dist: torch>=2.8
 Requires-Dist: torch_geometric
 Dynamic: author
 Dynamic: author-email