onnx-diagnostic 0.8.6__py3-none-any.whl → 0.8.8__py3-none-any.whl

onnx_diagnostic/export/dynamic_shapes.py

@@ -352,6 +352,19 @@ class CoupleInputsDynamicShapes:
                     else None
                 )
             assert type(inputs) is dict, f"Unexpected type for inputs {type(inputs)}"
+            if set(inputs) != set(ds):
+                not_in_ds = {k for k in inputs if k not in ds}
+                not_in_inputs = {k for k in ds if k not in inputs}
+                assert not_in_inputs == {"kwargs"} and set(ds["kwargs"]) == not_in_ds, (
+                    f"Keys mismatch between inputs {set(inputs)} and ds={set(ds)}, "
+                    f"inputs={string_type(inputs, with_shape=True)}, ds={ds}, "
+                    f"not_in_ds={not_in_ds}, not_in_inputs={not_in_inputs}"
+                )
+                # Tweak...
+                kws = ds["kwargs"]
+                del ds["kwargs"]
+                ds.update(kws)
+
             assert set(inputs) == set(ds), (
                 f"Keys mismatch between inputs {set(inputs)} and ds={set(ds)}, "
                 f"inputs={string_type(inputs, with_shape=True)}, ds={ds}"
@@ -366,13 +379,15 @@ class CoupleInputsDynamicShapes:
             return dvalue if dvalue else None
 
         # A custom class.
-        assert inputs.__class__ in torch.utils._pytree.SUPPORTED_NODES, (
+        assert inputs is None or inputs.__class__ in torch.utils._pytree.SUPPORTED_NODES, (
             f"Class {inputs.__class__.__name__!r} was not registered using "
             f"torch.utils._pytree.register_pytree_node, it is not possible to "
             f"map this class with the given dynamic shapes."
         )
         if flatten_unflatten:
             flatunflat = flatten_unflatten_for_dynamic_shapes(inputs)
+            if isinstance(flatunflat, (list, tuple, dict)) and len(flatunflat) == 0:
+                return flatunflat
             res = cls._generic_walker_step(
                 processor, flatunflat, ds, flatten_unflatten=flatten_unflatten
             )
@@ -667,6 +682,11 @@ class ModelInputs:
             if self.signature
             else None
         )
+        self.forward_parameters_kinds = (
+            {p.name: p.kind for p in self.signature.parameters.values()}
+            if self.signature
+            else None
+        )
         self.forward_ordered_parameter_names = (
             list(self.signature.parameters) if self.signature else None
         )
@@ -973,7 +993,13 @@ class ModelInputs:
             len(s1) == 1
         ), f"Different numbers of positional arguments {s1} for {self.full_name}"
         s2 = set(tuple(sorted(set(i[1]))) for i in self.inputs)
-        assert len(s2) == 1, f"Different named arguments {s2} for {self.full_name}"
+        assert len(s2) > 0, f"empty {s2} for {self.full_name}"
+        if len(s2) > 1:
+            # We need to keep the largest set of inputs, the one including all the others.
+            sum_s2 = set()
+            for s in s2:
+                sum_s2 |= set(s)
+            s2 = {tuple(sum_s2)}
         args = []
         kwargs = {}
         for i in range(s1.pop()):
@@ -993,12 +1019,18 @@ class ModelInputs:
                 f"\ninputs[1]={string_type(self.inputs[1], with_shape=True)}"
             )
 
-            objs = [_[1][name] for _ in self.inputs]
+            objs = [_[1][name] for _ in self.inputs if name in _[1]]
             kwargs[name] = self.guess_dynamic_shape_object(
                 *objs,
                 auto=auto if isinstance(auto, bool) else f"{auto}_{i}I",
                 msg=lambda name=name: f" failing input {name!r}",
             )
+        # reordering
+        if kwargs is not None and self.forward_ordered_parameter_names:
+            kwargs1 = {
+                p: kwargs[p] for p in self.forward_ordered_parameter_names if p in kwargs
+            }
+            kwargs = {**kwargs1, **{k: v for k, v in kwargs.items() if k not in kwargs1}}
         return tuple(args), kwargs
 
     def move_to_kwargs(
@@ -1061,6 +1093,16 @@ class ModelInputs:
                 f"and kwargs={set(kwargs)}, "
                 f"forward_ordered_parameter_names={self.forward_ordered_parameter_names}"
             )
+        if kwargs is not None and self.forward_ordered_parameter_names:
+            kwargs1 = {
+                p: kwargs[p] for p in self.forward_ordered_parameter_names if p in kwargs
+            }
+            kwargs = {**kwargs1, **{k: v for k, v in kwargs.items() if k not in kwargs1}}
+        if kw_dyn is not None and self.forward_ordered_parameter_names:
+            kw_dyn1 = {
+                p: kw_dyn[p] for p in self.forward_ordered_parameter_names if p in kw_dyn
+            }
+            kw_dyn = {**kw_dyn1, **{k: v for k, v in kw_dyn.items() if k not in kw_dyn1}}
         return args, kwargs, (tuple(), kw_dyn)
 
     def validate_inputs_for_export(

onnx_diagnostic/export/shape_helper.py

@@ -210,6 +210,7 @@ def make_fake_with_dynamic_dimensions(
     This uses function :func:`onnx_diagnostic.helpers.fake_tensor_helper.make_fake`.
     Parameter ``existing`` is used to reused the same object when the dynamic
     dimension is given the same name as another one.
+    This function works with caches only if ``transformers>=4.57``.
 
     A simple tensor:
 

onnx_diagnostic/ext_test_case.py

@@ -1267,6 +1267,7 @@ class ExtTestCase(unittest.TestCase):
             :class:`onnx_diagnostic.helpers.ort_session.InferenceSessionForTorch`
         """
         from .helpers import string_type, string_diff, max_diff
+        from .helpers.torch_helper import torch_deepcopy
         from .helpers.rt_helper import make_feeds
         from .helpers.ort_session import InferenceSessionForTorch
 
@@ -1283,6 +1284,12 @@ class ExtTestCase(unittest.TestCase):
                 model_file = proto
                 name = proto
                 proto = onnx.load(name)
+            elif hasattr(proto, "save"):
+                name = f"{test_name}.onnx"
+                proto.save(name)
+                proto = onnx.load(name)
+            elif hasattr(proto, "model_proto"):
+                proto = proto.model_proto
             elif not self.unit_test_going():
                 assert isinstance(
                     proto, onnx.ModelProto
@@ -1341,9 +1348,9 @@ class ExtTestCase(unittest.TestCase):
             if copy_inputs:
                 expected = [
                     (
-                        model(*copy.deepcopy(inp))
+                        model(*torch_deepcopy(inp))
                         if isinstance(inp, tuple)
-                        else model(**copy.deepcopy(inp))
+                        else model(**torch_deepcopy(inp))
                     )
                     for inp in inputs
                 ]

onnx_diagnostic/helpers/bench_run.py

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

onnx_diagnostic/helpers/cache_helper.py

@@ -28,14 +28,6 @@ class CacheKeyValue:
             ]
             self.key_cache = [layer.keys for layer in layers]
             self.value_cache = [layer.values for layer in layers]
-            if None in self.key_cache or None in self.value_cache:
-                from .helper import string_type
-
-                raise AssertionError(
-                    f"issue with key_cache={string_type(self.key_cache)}, "
-                    f"or value_cache={string_type(self.value_cache)}, "
-                    f"cache.layers={string_type(cache.layers)}"
-                )
         elif cache is not None and hasattr(cache, "key_cache"):
             self.key_cache = cache.key_cache
             self.value_cache = cache.value_cache

onnx_diagnostic/helpers/fake_tensor_helper.py

@@ -105,6 +105,8 @@ class FakeTensorContext:
         reduced_tensor = self.from_tensor(true_tensor, static_shapes=True).sum(
             axis=tuple(sorted(sh)), keepdim=True
         )
+        if len(reduced_tensor.shape) == 0 == len(new_shape):
+            return reduced_tensor
         return reduced_tensor.expand(*new_shape)
 
     def make_fake(self, x: Any) -> Optional["FakeTensor"]:  # noqa: F821
@@ -144,19 +146,22 @@ class FakeTensorContext:
         """
         See
         :func:`onnx_diagnostic.export.shape_helper.make_fake_with_dynamic_dimensions`.
+        If caches are used, it requires ``transformers>=4.57``.
         """
         if x is None:
             return None, None
-        if isinstance(x, (list, tuple)):
+        if type(x) in (list, tuple):
             return x.__class__(
                 [
                     self.make_fake_with_dynamic_dimensions(i, dynamic_shapes=ds)
                     for i, ds in zip(x, dynamic_shapes)
                 ]
             )
-        if isinstance(x, dict):
+        if type(x) is dict:
             return {
-                k: self.make_fake_with_dynamic_dimensions(v, dynamic_shapes=dynamic_shapes[k])
+                k: self.make_fake_with_dynamic_dimensions(
+                    v, dynamic_shapes=dynamic_shapes[k] if dynamic_shapes else None
+                )
                 for k, v in x.items()
             }
         if x.__class__.__name__ in {"DynamicCache", "StaticCache", "HybridCache"}:
@@ -187,6 +192,17 @@ class FakeTensorContext:
                 x.cross_attention_cache, dynamic_shapes=dynamic_shapes[1]
             )
             return x
+        if x.__class__.__name__ == "BaseModelOutput":
+            assert (
+                list(x.keys()) == ["last_hidden_state"] and x.last_hidden_state is not None
+            ), (
+                f"Field 'last_hidden_state' is empty for {type(x)} or other fields "
+                f"{list(x.keys())} are used."
+            )
+            x.last_hidden_state = self.make_fake_with_dynamic_dimensions(
+                x.last_hidden_state, dynamic_shapes=dynamic_shapes[0]
+            )
+            return x
         if hasattr(x, "shape"):
             assert dynamic_shapes is None or isinstance(dynamic_shapes, dict), (
                 f"dynamic_shapes must be a dictionary at this stage but "
@@ -197,9 +213,11 @@ class FakeTensorContext:
             for idim, dim in enumerate(x.shape):
                 if dynamic_shapes is not None and idim in dynamic_shapes:
                     s = dynamic_shapes[idim]
+                    if s.__class__.__name__ == "Dim":
+                        s = s.__name__
                     assert isinstance(s, str), (
                         f"Unexpected type {type(s)} in dynamic_shapes={dynamic_shapes} "
-                        f"at index {idim}"
+                        f"at index {idim}, self._mapping_str={self._mapping_str}"
                     )
                     if s in self._mapping_str:
                         dim = self._mapping_str[s]
@@ -217,10 +235,13 @@ class FakeTensorContext:
 
                 x = torch.empty(tuple(new_shape), dtype=x.dtype, device=x.device)
 
-            t = self.fake_reshape(x, dynamic_shapes)  # type: ignore[arg-type]
+            t = self.fake_reshape(x, dynamic_shapes) if dynamic_shapes else x  # type: ignore[arg-type]
             assert t.device == x.device, f"device mismatch {x.device} -> {t.device}"
             assert t.dtype == x.dtype, f"dtype mismatch {x.dtype} -> {t.dtype}"
             return t
+        if isinstance(x, (int, bool, float)):
+            # It is a constant, we don't change that.
+            return x
         from ..helpers import string_type
 
         raise TypeError(

onnx_diagnostic/helpers/helper.py

@@ -704,9 +704,35 @@ def string_type(
     if obj.__class__.__name__ == "VirtualTensor":
         if verbose:
             print(f"[string_type] TT4:{type(obj)}")
+
+        def _torch_sym_int_to_str(value: "torch.SymInt") -> Union[int, str]:  #  noqa: F821
+            if isinstance(value, str):
+                return value
+            if hasattr(value, "node") and isinstance(value.node, str):
+                return f"{value.node}"
+
+            from torch.fx.experimental.sym_node import SymNode
+
+            if hasattr(value, "node") and isinstance(value.node, SymNode):
+                # '_expr' is safer than expr
+                return str(value.node._expr).replace(" ", "")
+
+            try:
+                val_int = int(value)
+                return val_int
+            except (
+                TypeError,
+                ValueError,
+                AttributeError,
+                torch.fx.experimental.symbolic_shapes.GuardOnDataDependentSymNode,
+            ):
+                pass
+
+            raise AssertionError(f"Unable to convert {value!r} into string")
+
         return (
             f"{obj.__class__.__name__}(name={obj.name!r}, "
-            f"dtype={obj.dtype}, shape={obj.shape})"
+            f"dtype={obj.dtype}, shape={tuple(_torch_sym_int_to_str(_) for _ in obj.shape)})"
         )
 
     if obj.__class__.__name__ == "KeyValuesWrapper":
@@ -775,6 +801,9 @@ def string_type(
             print(f"[string_type] TT8:{type(obj)}")
         return repr(obj).replace(" ", "").replace("\n", " ")
 
+    if isinstance(obj, torch.fx.proxy.Proxy):
+        return repr(obj)
+
     if ignore:
         if verbose:
             print(f"[string_type] CACHE4:{type(obj)}")

onnx_diagnostic/helpers/log_helper.py

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

onnx_diagnostic/helpers/optim_helper.py

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

onnx_diagnostic/helpers/ort_session.py

@@ -1,3 +1,4 @@
+import os
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 import onnx
 import numpy as np
@@ -76,6 +77,10 @@ class _InferenceSession:
                     session_options.enable_profiling = enable_profiling
                 if optimized_model_filepath:
                     session_options.optimized_model_filepath = optimized_model_filepath
+                    session_options.add_session_config_entry(
+                        "session.optimized_model_external_initializers_file_name",
+                        f"{os.path.splitext(os.path.split(optimized_model_filepath)[-1])[0]}.data",
+                    )
                 if log_severity_level is not None:
                     session_options.log_severity_level = log_severity_level
                 if log_verbosity_level is not None:

onnx_diagnostic/tasks/image_text_to_text.py

@@ -13,6 +13,10 @@ from .data import get_data
 __TASK__ = "image-text-to-text"
 
 
+def should_have_vision_config(config):
+    return config.architectures != ["FuyuForCausalLM"]
+
+
 def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     kwargs: Dict[str, Any] = {}
@@ -168,10 +172,10 @@ def _get_inputs_gemma3(
         assert expected & set(
             dummies
         ), f"Unable to find expected inputs {expected} in loaded inputs {set(dummies)}"
-        assert sequence_length == dummies["input_ids"].shape[-1], (
-            f"sequence_length={sequence_length} != {dummies['input_ids'].shape[-1]} for "
-            f"model class {model.__class__.__name__}"
-        )
+        # assert sequence_length == dummies["input_ids"].shape[-1], (
+        #    f"sequence_length={sequence_length} != {dummies['input_ids'].shape[-1]} for "
+        #    f"model class {model.__class__.__name__}"
+        # )
         assert batch_size == dummies["input_ids"].shape[0], (
             f"batch_size={batch_size} != {dummies['input_ids'].shape[0]} for "
             f"model class {model.__class__.__name__}"
@@ -477,7 +481,8 @@ def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]:
                 "hidden_size",
                 "pad_token_id",
             )
-            check_hasattr(config, "vision_config", ("image_token_index", "image_token_id"))
+            if should_have_vision_config(config):
+                check_hasattr(config, "vision_config", ("image_token_index", "image_token_id"))
             text_config = True
         else:
             check_hasattr(
@@ -491,7 +496,8 @@ def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]:
                 "vision_config",
             )
             text_config = False
-        check_hasattr(config.vision_config, ("num_channels", "in_chans", "in_channels"))
+        if should_have_vision_config(config):
+            check_hasattr(config.vision_config, ("num_channels", "in_chans", "in_channels"))
     kwargs = dict(
         head_dim=(
             16
@@ -552,17 +558,21 @@ def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]:
         ),
         width=(
             224
-            if config is None or not hasattr(config.vision_config, "image_size")
+            if config is None
+            or not should_have_vision_config(config)
+            or not hasattr(config.vision_config, "image_size")
             else config.vision_config.image_size
         ),
         height=(
             224
-            if config is None or not hasattr(config.vision_config, "image_size")
+            if config is None
+            or not should_have_vision_config(config)
+            or not hasattr(config.vision_config, "image_size")
             else config.vision_config.image_size
         ),
         num_channels=(
             3
-            if config is None
+            if config is None or not should_have_vision_config(config)
             else _pick(config.vision_config, "num_channels", "in_chans", "in_channels")
         ),
         pad_token_id=(

onnx_diagnostic/tasks/text2text_generation.py

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

onnx_diagnostic/tasks/text_generation.py

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