onnx-diagnostic 0.7.4__py3-none-any.whl → 0.7.6__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])))
 
 
 #####################

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/hghub/model_inputs.py

@@ -144,6 +144,11 @@ def get_untrained_model_with_inputs(
                 f"[get_untrained_model_with_inputs] config._attn_implementation="
                 f"{config._attn_implementation!r}"  # type: ignore[union-attr]
             )
+    elif verbose:
+        print(
+            f"[get_untrained_model_with_inputs] default config._attn_implementation="
+            f"{getattr(config, '_attn_implementation', '?')!r}"  # type: ignore[union-attr]
+        )
 
     if type(config) is dict and "_diffusers_version" in config:
         import diffusers

onnx_diagnostic/torch_models/validate.py

@@ -288,6 +288,7 @@ def validate_model(
     repeat: int = 1,
     warmup: int = 0,
     inputs2: int = 1,
+    output_names: Optional[List[str]] = None,
 ) -> Tuple[Dict[str, Union[int, float, str]], Dict[str, Any]]:
     """
     Validates a model.
@@ -338,6 +339,7 @@ def validate_model(
     :param inputs2: checks that the second set of inputs is reunning as well,
         this ensures that the model does support dynamism, the value is used
         as an increment to the first set of values (added to dimensions)
+    :param output_names: output names the onnx exporter should use
     :return: two dictionaries, one with some metrics,
         another one with whatever the function produces
 
@@ -433,6 +435,7 @@ def validate_model(
         )
         print(f"[validate_model] exporter={exporter!r}, optimization={optimization!r}")
         print(f"[validate_model] dump_folder={dump_folder!r}")
+        print(f"[validate_model] output_names={output_names}")
         summary["model_id"] = model_id
         summary["model_subfolder"] = subfolder or ""
 
@@ -631,6 +634,7 @@ def validate_model(
                     optimization=optimization,
                     do_run=do_run,
                     dump_folder=dump_folder,
+                    output_names=output_names,
                 )
         else:
             data["inputs_export"] = data["inputs"]
@@ -643,6 +647,7 @@ def validate_model(
                 optimization=optimization,
                 do_run=do_run,
                 dump_folder=dump_folder,
+                output_names=output_names,
             )
         summary.update(summary_export)
 
@@ -868,6 +873,7 @@ def call_exporter(
     optimization: Optional[str] = None,
     do_run: bool = False,
     dump_folder: Optional[str] = None,
+    output_names: Optional[List[str]] = None,
 ) -> Tuple[Dict[str, Union[int, float, str]], Dict[str, Any]]:
     """
     Calls an exporter on a model;
@@ -880,6 +886,7 @@ def call_exporter(
     :param optimization: optimization to do
     :param do_run: runs and compute discrepancies
     :param dump_folder: to dump additional information
+    :param output_names: list of output names to use with the onnx exporter
     :return: two dictionaries, one with some metrics,
         another one with whatever the function produces
     """
@@ -902,6 +909,7 @@ def call_exporter(
             quiet=quiet,
             verbose=verbose,
             optimization=optimization,
+            output_names=output_names,
         )
         return summary, data
     if exporter == "custom" or exporter.startswith("custom"):
@@ -913,6 +921,7 @@ def call_exporter(
             verbose=verbose,
             optimization=optimization,
             dump_folder=dump_folder,
+            output_names=output_names,
         )
         return summary, data
     if exporter == "modelbuilder":
@@ -923,6 +932,7 @@ def call_exporter(
             quiet=quiet,
             verbose=verbose,
             optimization=optimization,
+            output_names=output_names,
         )
         return summary, data
     raise NotImplementedError(
@@ -1090,7 +1100,7 @@ def validate_onnx_model(
     """
     import onnxruntime
 
-    def _mk(key):
+    def _mk(key, flavour=flavour):
         return f"{key}_{flavour}" if flavour else key
 
     summary: Dict[str, Any] = {}
@@ -1145,7 +1155,7 @@ def validate_onnx_model(
     )
     sess = _quiet_or_not_quiet(
         quiet,
-        _mk("onnx_ort_create"),
+        _mk("create_onnx_ort"),
         summary,
         data,
         (lambda source=source, providers=providers: cls_runtime(source, providers)),
@@ -1180,7 +1190,7 @@ def validate_onnx_model(
 
         got = _quiet_or_not_quiet(
             quiet,
-            _mk(f"time_onnx_ort_run{suffix}"),
+            _mk(f"run_onnx_ort{suffix}"),
             summary,
             data,
             (lambda sess=sess, feeds=feeds: sess.run(None, feeds)),
@@ -1211,6 +1221,7 @@ def call_torch_export_onnx(
     quiet: bool = False,
     verbose: int = 0,
     optimization: Optional[str] = None,
+    output_names: Optional[List[str]] = None,
 ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
     """
     Exports a model into onnx.
@@ -1222,6 +1233,7 @@ def call_torch_export_onnx(
     :param quiet: catch exception or not
     :param verbose: verbosity
     :param optimization: optimization to do
+    :param output_names: output names to use
     :return: two dictionaries, one with some metrics,
         another one with whatever the function produces
     """
@@ -1276,6 +1288,8 @@ def call_torch_export_onnx(
             print("[call_torch_export_onnx] dynamo=False so...")
             print(f"[call_torch_export_onnx] args={string_type(args, with_shape=True)}")
             print(f"[call_torch_export_onnx] kwargs={string_type(kwargs, with_shape=True)}")
+    if output_names:
+        export_export_kwargs["output_names"] = output_names
     if opset:
         export_export_kwargs["opset_version"] = opset
     if verbose:
@@ -1346,6 +1360,7 @@ def call_torch_export_model_builder(
     quiet: bool = False,
     verbose: int = 0,
     optimization: Optional[str] = None,
+    output_names: Optional[List[str]] = None,
 ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
     """
     Exports a model into onnx with :epkg:`ModelBuilder`.
@@ -1356,6 +1371,7 @@ def call_torch_export_model_builder(
     :param quiet: catch exception or not
     :param verbose: verbosity
     :param optimization: optimization to do
+    :param output_names: list of output names to use
     :return: two dictionaries, one with some metrics,
         another one with whatever the function produces
     """
@@ -1369,6 +1385,9 @@ def call_torch_export_model_builder(
     provider = data.get("model_device", "cpu")
     dump_folder = data.get("model_dump_folder", "")
     assert dump_folder, "dump_folder cannot be empty with ModelBuilder"
+    assert (
+        not output_names
+    ), f"output_names not empty, not supported yet, output_names={output_names}"
     cache_dir = os.path.join(dump_folder, "cache_mb")
     if not os.path.exists(cache_dir):
         os.makedirs(cache_dir)
@@ -1408,6 +1427,7 @@ def call_torch_export_custom(
     verbose: int = 0,
     optimization: Optional[str] = None,
     dump_folder: Optional[str] = None,
+    output_names: Optional[List[str]] = None,
 ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
     """
     Exports a model into onnx.
@@ -1420,6 +1440,7 @@ def call_torch_export_custom(
     :param verbose: verbosity
     :param optimization: optimization to do
     :param dump_folder: to store additional information
+    :param output_names: list of output names to use
     :return: two dictionaries, one with some metrics,
         another one with whatever the function produces
     """
@@ -1504,6 +1525,8 @@ def call_torch_export_custom(
     )
     if opset:
         kws["target_opset"] = opset
+    if output_names:
+        kws["output_names"] = output_names
 
     epo, opt_stats = _quiet_or_not_quiet(
         quiet,

{onnx_diagnostic-0.7.4.dist-info → onnx_diagnostic-0.7.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: onnx-diagnostic
-Version: 0.7.4
+Version: 0.7.6
 Summary: Investigate ONNX models
 Home-page: https://github.com/sdpython/onnx-diagnostic
 Author: Xavier Dupré