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

onnx_diagnostic/export/api.py

@@ -1,11 +1,18 @@
 import inspect
 import os
 import textwrap
+import time
+from collections.abc import Mapping, Iterable
 from typing import Any, Callable, Dict, List, Optional, Sequence, Set, Tuple, Union
 import torch
 from .dynamic_shapes import ModelInputs
 from .onnx_plug import EagerDirectReplacementWithOnnx
-from ..helpers import string_type
+from ..helpers import flatten_object, max_diff, string_diff, string_type
+from ..helpers.cache_helper import CacheKeyValue
+from ..helpers.torch_helper import torch_deepcopy
+from ..helpers.rt_helper import make_feeds
+from ..helpers.onnx_helper import pretty_onnx
+from ..reference import OnnxruntimeEvaluator
 
 
 def get_main_dispatcher(
@@ -314,10 +321,11 @@ def to_onnx(
     raise ValueError(f"Unknown exporter={exporter!r}")
 
 
-class _WrapperToExportMethodToOnnx(torch.nn.Module):
+class WrapperToExportMethodToOnnx(torch.nn.Module):
     """
     Wraps an existing models in order to spy on inputs.
-    This is used by :func:`onnx_diagnostic.export.api.method_to_onnx`.
+    This is used by :func:`onnx_diagnostic.export.api.method_to_onnx`
+    or :ref:`l-plot-tiny-llm-export-method-generate` for an example.
     """
 
     def __init__(
@@ -342,6 +350,8 @@ class _WrapperToExportMethodToOnnx(torch.nn.Module):
         patch_kwargs: Optional[Dict[str, Any]] = None,
         skip_kwargs_names: Optional[Set[str]] = None,
         dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
+        dynamic_batch_for: Optional[Sequence[Union[int, str]]] = None,
+        expand_batch_for: Optional[Sequence[Union[int, str]]] = None,
     ):
         super().__init__()
         self._model_to_call = mod
@@ -351,13 +361,17 @@ class _WrapperToExportMethodToOnnx(torch.nn.Module):
             if method_name == "forward"
             else getattr(mod, method_name)
         )
+        self._signature = inspect.signature(self._method_call)
         self._inputs: List[Tuple[Tuple[Any, ...], Dict[str, Any]]] = []
+        self._outputs: List[Any] = []
         self._convert_after_n_calls = convert_after_n_calls
         self._patch_kwargs = patch_kwargs
         self._method_src = None
         self.verbose = verbose
         self.skip_kwargs_names = skip_kwargs_names
         self.dynamic_shapes = dynamic_shapes
+        self.expand_batch_for = expand_batch_for
+        self.dynamic_batch_for = dynamic_batch_for
         self._to_onnx_kwargs = dict(
             input_names=input_names,
             target_opset=target_opset,
@@ -375,42 +389,117 @@ class _WrapperToExportMethodToOnnx(torch.nn.Module):
             inline=inline,
         )
         self._export_done = False
+        self._serialization_classes: Set[type] = set()
 
     def __str__(self) -> str:
+        "usual"
         return self.__repr__()
 
     def __repr__(self) -> str:
+        "usual"
         return (
             f"{self.__class__.__name__}({self._model_to_call.__class__.__name__}."
             f"{self._method_name})"
         )
 
+    def _collect_classes(self, obj):
+        if obj is None or isinstance(obj, torch.Tensor):
+            return
+        cls = type(obj)
+        if cls.__module__ not in ("builtins",):
+            self._serialization_classes.add(cls)
+        if hasattr(obj, "__dict__"):
+            for v in vars(obj).values():
+                self._collect_classes(v)
+            return
+        if isinstance(obj, Mapping):
+            for v in obj.values():
+                self._collect_classes(v)
+            return
+        if isinstance(obj, Iterable) and not isinstance(obj, (str, bytes)):
+            for v in obj:
+                self._collect_classes(v)
+            return
+
+    def _reorder_kwargs(self, kwargs):
+        new_kwargs = {k: kwargs[k] for k in self._signature.parameters if k in kwargs}
+        for k, v in kwargs.items():
+            if k not in new_kwargs:
+                new_kwargs[k] = v
+        return new_kwargs
+
     def forward(self, *args, **kwargs):
         if not self._export_done:
-            self._inputs.append(
-                (
-                    args,
-                    (
-                        kwargs
-                        if not kwargs or not self.skip_kwargs_names
-                        else {
-                            k: v for k, v in kwargs.items() if k not in self.skip_kwargs_names
-                        }
-                    ),
-                )
+            inp_args = args
+            # filters out the inputs not desired, int, float, bool, None
+            # are considered as constant for the exporter, they are removed
+            # from the named arguments.
+            inp_kwargs = (
+                kwargs
+                if not kwargs
+                else {
+                    k: v
+                    for k, v in kwargs.items()
+                    if v is not None
+                    and (not self.skip_kwargs_names or k not in self.skip_kwargs_names)
+                    and not isinstance(v, (bool, int, float))
+                }
             )
+            inp_args, inp_kwargs = torch_deepcopy((inp_args, inp_kwargs))
+            # reorders the parameter following the method signature.
+            inp_kwargs = self._reorder_kwargs(inp_kwargs)
+            # stores the inputs
+            self._inputs.append((inp_args, inp_kwargs))
+
             if self.verbose:
                 print(
                     f"[method_to_onnx] input[{len(self._inputs)-1}]: "
                     f"{string_type(self._inputs[-1], with_shape=True)}"
                 )
+
             if len(self._inputs) >= self._convert_after_n_calls:
+                # conversion starts after _convert_after_n_calls calls to the forward method
+                name = os.path.splitext(self._to_onnx_kwargs["filename"])[0]
+                input_file = f"{name}.inputs.pt"
+                self._input_file = input_file
+                if self.verbose:
+                    print(
+                        f"[method_to_onnx] save {len(self._inputs)} inputs in {input_file!r}"
+                    )
+                torch.save(self._inputs, input_file)
                 self._convert_method_to_onnx()
-                del self._inputs[:]
                 self._export_done = True
-        return self._method_call(*args, **kwargs)
+
+        # calls the inner method (no change here)
+        begin = time.perf_counter()
+        res = self._method_call(*args, **kwargs)
+        duration = time.perf_counter() - begin
+        self._collect_classes([args, kwargs, res])
+        if self._inputs:
+            # stores the outputs if discrepancies need to be checked
+            self._outputs.append((torch_deepcopy(res), duration))
+            assert len(self._inputs) == len(self._outputs), (
+                f"Number of inputs {len(self._inputs)} and "
+                f"outputs {len(self._outputs)} are different."
+            )
+            if self._export_done:
+                name = os.path.splitext(self._to_onnx_kwargs["filename"])[0]
+                output_file = f"{name}.outputs.pt"
+                if self.verbose:
+                    print(
+                        f"[method_to_onnx] save {len(self._outputs)} "
+                        f"outputs in {output_file!r}"
+                    )
+                torch.save(self._outputs, output_file)
+                self._output_file = output_file
+                del self._inputs[:]
+                del self._outputs[:]
+        return res
 
     def _convert_method_to_onnx(self):
+        for args, kwargs in self._inputs:
+            self._serialization_classes |= {type(a) for a in args}
+            self._serialization_classes |= {type(a) for a in kwargs.values()}
 
         def make_method(self):
             inner_sig = inspect.signature(self._method_call)
@@ -450,9 +539,24 @@ class _WrapperToExportMethodToOnnx(torch.nn.Module):
             if self.verbose:
                 print(f"[method_to_onnx] guess_dynamic_shapes={string_type(ds)}")
             a, kw, nds = mi.move_to_kwargs(*self._inputs[-1], ds)
+            if self.dynamic_batch_for:
+                nds = (
+                    self._dynamic_batch_dimension(nds[0], self.dynamic_batch_for),
+                    self.rename_dynamic_shapes(
+                        self._dynamic_batch_dimension(nds[1], self.dynamic_batch_for),
+                        verbose=self.verbose,
+                    ),
+                )
+                if self.verbose:
+                    print(f"[method_to_onnx] dynamic_batch_for={self.dynamic_batch_for}")
+                    print(f"[method_to_onnx] dynamic_shapes with batch={nds}")
         else:
             a, kw = self._inputs[-1]
             nds = [self.dynamic_shapes]
+        if self.expand_batch_for:
+            # extends the inputs to artificially create a batch dimension != 1.
+            a = self._expand_batch_dimension(a, self.expand_batch_for)
+            kw = self._expand_batch_dimension(kw, self.expand_batch_for)
         if self.verbose:
             print(f"[method_to_onnx] export args={string_type(a, with_shape=True)}")
             print(f"[method_to_onnx] export kwargs={string_type(kw, with_shape=True)}")
@@ -477,6 +581,365 @@ class _WrapperToExportMethodToOnnx(torch.nn.Module):
                 **self._to_onnx_kwargs,
             )
 
+    @classmethod
+    def make_empty_cache_from_others(cls, examples: List[Any]) -> Any:
+        """Builds an empty cache based on existing one."""
+        unique_types = {type(t) for t in examples}
+        assert (
+            len(unique_types) == 1
+        ), f"Unable to guess an empty cache from {string_type(examples, with_shape=True)}"
+        unique_type = unique_types.pop()
+        if unique_type == torch.Tensor:
+            shapes = [t.shape for t in examples]
+            assert len(set(shapes)) > 1, f"Unable to guess an empty shape from shapes {shapes}"
+            ranks = {len(s) for s in shapes}
+            assert len(ranks) == 1, f"Ranks are different in {shapes}"
+            rank = ranks.pop()
+            new_shape = []
+            for i in range(rank):
+                dims = [t.shape[i] for t in examples]
+                if len(set(dims)) == 1:
+                    new_shape.append(dims[0])
+                else:
+                    # The empty shape
+                    new_shape.append(0)
+            example = examples[0]
+            return torch.empty(tuple(new_shape), dtype=example.dtype, device=example.device)
+        assert (
+            unique_type.__name__ == "DynamicCache"
+        ), f"This is not implemented for class {unique_type}"
+        caches = [CacheKeyValue(dc) for dc in examples]
+        caches_list = [dc.aslist() for dc in caches]
+        empty = [
+            cls.make_empty_cache_from_others([caches_list[i][k] for i in range(len(examples))])
+            for k in range(len(caches_list[0]))
+        ]
+        empty_cache = CacheKeyValue(
+            empty, cls_layers=caches[0].cls_layers
+        ).make_dynamic_cache()
+        return empty_cache
+
+    @classmethod
+    def add_empty_cache_if_needed(cls, inputs: List[Any]) -> List[Any]:
+        """
+        Adds empty cache if needed as onnxruntime needs an empty cache,
+        not a missing cache. It only works if inputs are defined as a dictionary.
+        """
+        if all(isinstance(t, tuple) for t in inputs) and all(
+            len(t) == 2 and isinstance(t[0], tuple) and isinstance(t[1], dict) and not t[0]
+            for t in inputs
+        ):
+            dict_part = [t[1] for t in inputs]
+            res = cls.add_empty_cache_if_needed(dict_part)
+            return [(tuple(), d) for d in res]
+        if any(not isinstance(t, dict) for t in inputs):
+            return inputs
+        all_keys = set()
+        for input_set in inputs:
+            all_keys |= set(input_set)
+        # even though the inputs are defined as a dictionary, it is better
+        # to keep the same order
+        ordered = None
+        for input_set in inputs:
+            if set(input_set) == all_keys:
+                ordered = list(input_set)
+                break
+        new_inputs = []
+        for input_set in inputs:
+            if set(input_set) == all_keys:
+                new_inputs.append(input_set)
+                continue
+            missing = {k for k in all_keys if k not in input_set}
+            input_set_copy = input_set.copy()
+            for miss in missing:
+                input_set_copy[miss] = cls.make_empty_cache_from_others(
+                    [sub[miss] for sub in inputs if miss in sub]
+                )
+            new_inputs.append({k: input_set_copy[k] for k in ordered})  # type: ignore[union-attr]
+        return new_inputs
+
+    @classmethod
+    def _expand_batch_dimension(cls, obj: Any, expand_for: Sequence[Union[int, str]]) -> Any:
+        expand_for_args = {i for i in expand_for if isinstance(i, int)}
+        expand_for_kwargs = {i for i in expand_for if isinstance(i, str)}
+        if isinstance(obj, tuple):
+            return tuple(
+                o if i not in expand_for_args else cls._expand_batch_dimension_input(o, i)
+                for i, o in enumerate(obj)
+            )
+        assert isinstance(obj, dict), f"Unexpected type {type(obj)}"
+        return {
+            k: v if k not in expand_for_kwargs else cls._expand_batch_dimension_input(v, k)
+            for k, v in obj.items()
+        }
+
+    @classmethod
+    def _expand_batch_dimension_input(cls, obj: Any, msg: Union[str, int]) -> Any:
+        if isinstance(obj, torch.Tensor):
+            assert obj.shape[0] == 1, (
+                f"Are you sure to expoand input {msg!r}, "
+                f"batch size is not 1 and shape={obj.shape}"
+            )
+            sizes = [2, *obj.shape[1:]]
+            return obj.expand(*sizes)
+        if isinstance(obj, list):
+            return [
+                cls._expand_batch_dimension_input(o, f"{msg}[{i}]") for i, o in enumerate(obj)
+            ]
+        if obj.__class__.__name__ == "DynamicCache":
+            dc = CacheKeyValue(obj)
+            flat = dc.aslist()
+            flat = cls._expand_batch_dimension_input(flat, msg)
+            return CacheKeyValue(flat, cls_layers=dc.cls_layers).make_dynamic_cache()
+        # This might end up in an infinite loop if no registration is done.
+        flat, _spec = torch.utils._pytree.tree_flatten(obj)
+        assert (
+            not isinstance(flat, list) or len(flat) != 1 or type(flat[0]) is not type(obj)
+        ), f"class {type(obj)} was is not registered for serialization."
+        flat = cls._expand_batch_dimension_input(flat, msg)
+        return torch.utils._pytree.tree_unflatten(flat, _spec)
+
+    @classmethod
+    def _dynamic_batch_dimension(
+        cls, ds: Union[Tuple[Any, ...], Dict[str, Any]], dynamic_for: Sequence[Union[int, str]]
+    ) -> Union[Tuple[Any, ...], Dict[str, Any]]:
+        if isinstance(ds, tuple):
+            return tuple(
+                (v if i not in dynamic_for else cls._dynamic_batch_dimension_input(v, i))
+                for i, v in enumerate(ds)
+            )
+        return {
+            k: (v if k not in dynamic_for else cls._dynamic_batch_dimension_input(v, k))
+            for k, v in ds.items()
+        }
+
+    @classmethod
+    def _dynamic_batch_dimension_input(cls, ds: Any, msg: Union[str, int]) -> Any:
+        if isinstance(ds, dict) and all(isinstance(k, int) for k in ds):
+            ds[0] = "batch"
+            return {k: v for k, v in sorted(ds.items())}  # noqa: C416
+        if isinstance(ds, list):
+            return [
+                cls._dynamic_batch_dimension_input(o, f"{msg}[{i}]") for i, o in enumerate(ds)
+            ]
+        raise NotImplementedError(f"cannot make first dimension dynamic for batch for {ds}")
+
+    def check_discrepancies(
+        self, atol: float = 1e-4, rtol: float = 0.1, hist=(0.1, 0.01), verbose: int = 0
+    ) -> List[Dict[str, Union[str, int, float]]]:
+        """
+        Computes the discrepancies between the saved inputs and outputs
+        with the saved onnx model.
+
+        :param atol: absolute tolerance, recommended values, 1e-4 for float, 1e-2 flot float16
+        :param rtol: relative tolerance
+        :param hist: thresholds, the function determines the number of discrepancies
+            above that threshold.
+        :param verbose: verbosity
+        :return: results, a list of dictionaries, ready to be consumed by a dataframe
+        """
+        assert (
+            self._export_done
+        ), f"The onnx export was not done, only {len(self._inputs)} were stored."
+        assert os.path.exists(self._input_file), f"input file {self._input_file!r} not found"
+        assert os.path.exists(
+            self._output_file
+        ), f"output file {self._output_file!r} not found"
+        filename = self._to_onnx_kwargs["filename"]
+        assert isinstance(filename, str) and os.path.exists(
+            filename
+        ), f"onnx file {filename!r} not found"
+        classes = [
+            cls
+            for cls in self._serialization_classes
+            if cls
+            not in {
+                int,
+                float,
+                bool,
+                str,
+                torch.Tensor,
+                list,
+                set,
+                dict,
+                torch.device,
+                torch.dtype,
+            }
+        ]
+        if verbose:
+            print(f"[method_to_onnx.check_discrepancies] register classes {classes}")
+            print(f"[method_to_onnx.check_discrepancies] load {self._input_file!r}")
+        with torch.serialization.safe_globals(classes):
+            inputs = torch.load(self._input_file, weights_only=False)
+        if verbose:
+            print(f"[method_to_onnx.check_discrepancies] load {self._output_file!r}")
+        with torch.serialization.safe_globals(classes):
+            outputs = torch.load(self._output_file, weights_only=False)
+        assert len(inputs) == len(outputs), (
+            f"Unexpected number of inputs {len(inputs)} and outputs {len(outputs)}, "
+            f"inputs={string_type(inputs, with_shape=True)}, "
+            f"outputs={string_type(outputs, with_shape=True)}"
+        )
+        if verbose:
+            print(f"[method_to_onnx.check_discrepancies] create onnx session {filename!r}")
+        sess = OnnxruntimeEvaluator(filename, whole=True)
+        input_names = sess.input_names
+        if verbose:
+            print(f"[method_to_onnx.check_discrepancies] input_names={input_names}")
+            print(
+                f"[method_to_onnx.check_discrepancies] onnx_shapes="
+                f"{', '.join(pretty_onnx(i) for i in sess.input_types)}"
+            )
+        data = []
+        for i, (input, (output, latency)) in enumerate(
+            zip(self.add_empty_cache_if_needed(inputs), outputs)
+        ):
+            if verbose:
+                if verbose > 1:
+                    print(
+                        f"[method_to_onnx.check_discrepancies] process input {i}: "
+                        f"{string_type(input, with_shape=True)}"
+                    )
+                    print(
+                        f"[method_to_onnx.check_discrepancies] expects: "
+                        f"{string_type(output, with_shape=True)}"
+                    )
+                else:
+                    print(
+                        f"[method_to_onnx.check_discrepancies] process input {i} "
+                        f"#args={len(input[0])} #kwargs={len(input[1])}"
+                    )
+
+            flat_inputs = flatten_object(input, drop_keys=True)
+            if verbose > 1:
+                print(
+                    f"[method_to_onnx.check_discrepancies] "
+                    f"input={string_type(input, with_shape=True)}"
+                )
+                print(
+                    f"[method_to_onnx.check_discrepancies] "
+                    f"flat_inputs={string_type(flat_inputs, with_shape=True)}"
+                )
+            if len(flat_inputs) < len(input_names):
+                # not implemented yet, it is caused by a missing cache,
+                # which requires an empty cache instead
+                data.append(dict(index=i, duration_torch=latency, n_inputs=len(flat_inputs)))
+                continue
+            assert len(flat_inputs) == len(input_names), (
+                f"Length mismatch, expecting {len(input_names)} onnx inputs and got "
+                f"{len(flat_inputs)} flat torch inputs"
+            )
+            feeds = make_feeds(input_names, flat_inputs)
+            if verbose > 1:
+                print(
+                    f"[method_to_onnx.check_discrepancies] "
+                    f"feeds={string_type(feeds, with_shape=True)}"
+                )
+            begin = time.perf_counter()
+            ort_outputs = sess.run(None, feeds)
+            duration = time.perf_counter() - begin
+            diff = max_diff(output, ort_outputs, hist=hist)
+            if "rep" in diff and isinstance(diff["rep"], dict):
+                diff.update(diff["rep"])
+                del diff["rep"]
+            diff["SUCCESS"] = (
+                isinstance(diff["abs"], float)
+                and isinstance(diff["rel"], float)
+                and diff["abs"] < atol
+                and diff["rel"] < rtol
+            )
+            diff.update(
+                dict(
+                    index=i,
+                    duration_torch=latency,
+                    ort_duration=duration,
+                    n_inputs=len(flat_inputs),
+                )
+            )
+            if verbose > 1:
+                print(
+                    f"[method_to_onnx.check_discrepancies] ort output "
+                    f"{string_type(ort_outputs, with_shape=True)}"
+                )
+                print(f"[method_to_onnx.check_discrepancies] diff {string_diff(diff)}")
+            data.append(diff)
+        if verbose:
+            print("[method_to_onnx.check_discrepancies] done")
+        return data
+
+    @classmethod
+    def _apply_known_shape_pattern(
+        cls, shape: Dict[int, Any], pattern: Dict[int, str]
+    ) -> Dict[int, Any]:
+        return {k: pattern.get(k, v) for k, v in shape.items()}
+
+    @classmethod
+    def get_dynamic_shape_patterns(cls) -> Dict[str, Any]:
+        """
+        Returns the known patterns for the dynamic shapes.
+
+        .. runpython::
+            :showcode:
+
+            import pprint
+            from onnx_diagnostic.export.api import WrapperToExportMethodToOnnx
+            pprint.pprint(WrapperToExportMethodToOnnx.get_dynamic_shape_patterns())
+        """
+        return {
+            "LLM.text": {
+                "cache_position": {0: "seqlength"},
+                "past_key_values": {0: "batch", 2: "pastlength"},
+                "input_ids": {0: "batch", 1: "seqlength"},
+                "attention_mask": {0: "batch", 1: "totallength"},  # pastlength+seqlength
+            }
+        }
+
+    @classmethod
+    def rename_dynamic_shapes(cls, ds: Dict[str, Any], verbose: int = 0) -> Dict[str, Any]:
+        """
+        Renames the dynamic shapes with names.
+        Tries to rename any dynamic dimnesion dimension
+        before export. It is not very clever, it just tries
+        to recognize a known configuration based on input names.
+        Dimension names in dynamic shapes are renamed if *ds* has
+        the same number of named arguments as the one of the patterns
+        returned by function :meth:`get_dynamic_shape_patterns
+        <onnx_diagnostic.export.api.WrapperToExportMethodToOnnx.get_dynamic_shape_patterns>`.
+        """
+        is_shape = lambda s: isinstance(s, dict) and all(  # noqa: E731
+            isinstance(_, int) for _ in s
+        )
+        llm_patterns = cls.get_dynamic_shape_patterns()
+        for pattern_name, pattern_shape in llm_patterns.items():
+            if len(set(ds) & set(pattern_shape)) == len(pattern_shape):
+                if verbose:
+                    print(
+                        f"[method_to_onnx.rename_dynamic_shapes] "
+                        f"apply pattern shapes {pattern_name!r}"
+                    )
+                new_ds = {}
+                for k, v in ds.items():
+                    if k not in pattern_shape:
+                        new_ds[k] = v
+                        continue
+                    if is_shape(v):
+                        # A shape
+                        new_ds[k] = cls._apply_known_shape_pattern(v, pattern_shape[k])
+                    elif isinstance(v, list):
+                        # A cache
+                        new_ds[k] = [
+                            (
+                                cls._apply_known_shape_pattern(s, pattern_shape[k])
+                                if is_shape(s)
+                                else s
+                            )
+                            for s in v
+                        ]
+                return new_ds
+
+        # unchanged
+        return ds
+
 
 def method_to_onnx(
     mod: "torch.nn.Module",
@@ -499,6 +962,8 @@ def method_to_onnx(
     patch_kwargs: Optional[Dict[str, Any]] = None,
     skip_kwargs_names: Optional[Set[str]] = None,
     dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
+    dynamic_batch_for: Optional[Sequence[Union[int, str]]] = None,
+    expand_batch_for: Optional[Sequence[Union[int, str]]] = None,
 ) -> Callable:
     """
     Exports one method into ONNX for a module into ONNX.
@@ -528,12 +993,20 @@ def method_to_onnx(
     :param skip_kwargs_names: use default values for these parameters part of
         the signature of the method to export
     :param dynamic_shapes: dynamic shapes to use if the guessed ones are not right
+    :param dynamic_batch_for: LLM are usually called with a batch size equal to 1,
+        but the export may benefit from having a dynamic batch size,
+        this parameter forces the input specified in this set to have the first dimension
+        be dynamic
+    :param expand_batch_for: LLM are usually called with a batch size equal to 1,
+        but the export may benefit from having another value for the batch size,
+        this parameter forces the input specified in this set to be expanded
+        to 2 if the batch size is one
     :return: the output of the selected exporter, usually a structure including
         an onnx model
 
     See :ref:`l-plot-tiny-llm-export-method-generate` for an example.
     """
-    wrapped_model = _WrapperToExportMethodToOnnx(
+    wrapped_model = WrapperToExportMethodToOnnx(
         mod=mod,
         method_name=method_name,
         input_names=input_names,
@@ -554,5 +1027,7 @@ def method_to_onnx(
         patch_kwargs=patch_kwargs,
         skip_kwargs_names=skip_kwargs_names,
         dynamic_shapes=dynamic_shapes,
+        dynamic_batch_for=dynamic_batch_for,
+        expand_batch_for=expand_batch_for,
     )
     return wrapped_model

onnx_diagnostic/export/dynamic_shapes.py

@@ -329,7 +329,7 @@ class CoupleInputsDynamicShapes:
         if type(inputs) in (tuple, list, dict):
             # Type must be strict, some custom classes can inherit from those.
             assert type(inputs) is type(ds), (
-                f"Input type and dynamic shape type mush match but "
+                f"Input type and dynamic shapes type mush match but "
                 f"type(inputs)={type(inputs)}, type(ds)={type(ds)}, "
                 f"inputs={string_type(inputs, with_shape=True)}, ds={ds}"
             )
@@ -967,6 +967,8 @@ class ModelInputs:
         """
         Guesses the dynamic shapes for that module from two execution.
         If there is only one execution, then that would be static dimensions.
+        If the model signature is available, the kwargs are reordered following
+        the signature order, otherwise it follows the order given in the inputs.
 
         :param auto: if auto is True, use ``torch.export.Dim.AUTO`` for any
             dimension if the number of inputs is one,
@@ -1026,11 +1028,24 @@ class ModelInputs:
                 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}}
+        if kwargs:
+            if 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}}
+            else:
+                # We reorder the same the way the input were given.
+                use = None
+                params = set(kwargs)
+                for _args, kws in self.inputs:
+                    if set(kws) == params:
+                        use = kws
+                        break
+                if use:
+                    ordered = list(use)
+                    kwargs = {k: kwargs[k] for k in ordered}
+
         return tuple(args), kwargs
 
     def move_to_kwargs(

onnx_diagnostic/export/shape_helper.py

@@ -47,7 +47,6 @@ def all_dynamic_shapes_from_inputs(inputs: Any, dim_prefix: Any = "d") -> Any:
             make_dynamic_cache,
             make_encoder_decoder_cache,
             make_mamba_cache,
-            make_sliding_window_cache,
             make_static_cache,
         )
         from onnx_diagnostic.export.shape_helper import all_dynamic_shapes_from_inputs
@@ -77,13 +76,6 @@ def all_dynamic_shapes_from_inputs(inputs: Any, dim_prefix: Any = "d") -> Any:
                     ]
                 ),
             ),
-            make_sliding_window_cache(
-                [
-                    (torch.rand((4, 5, 6, 7)), torch.rand((4, 5, 6, 7))),
-                    (torch.rand((4, 5, 6, 7)), torch.rand((4, 5, 6, 7))),
-                    (torch.rand((4, 5, 6, 7)), torch.rand((4, 5, 6, 7))),
-                ]
-            ),
             make_static_cache(
                 [
                     (torch.rand((4, 5, 6, 7)), torch.rand((4, 5, 6, 7))),