fix

Author: xadupre

_doc/examples/plot_export_tiny_llm_method_generate.py

@@ -68,6 +68,7 @@ def generate_text(
     patch_kwargs=dict(patch_transformers=True),
     verbose=1,
     convert_after_n_calls=3,
+    skip_kwargs_names={"kwargs", "use_cache", "return_dict"},
 )
 
 # %%

_unittests/ut_export/test_api.py

@@ -134,7 +134,7 @@ def forward(self, x, y):
         self.assertExists(filename)
         src = method_to_call._method_src
         self.assertIn("f(self, x, y):", src)
-        self.assertIn("return self._call(x=x, y=y)", src)
+        self.assertIn("return self._method_call(x=x, y=y)", src)
         self.assertEqual(len(list(method_to_call.named_modules())), 2)
         sess = self.check_ort(filename)
         input_names = [i.name for i in sess.get_inputs()]
@@ -163,7 +163,7 @@ def forward(self, x=None, y=None):
         self.assertExists(filename)
         src = method_to_call._method_src
         self.assertIn("f(self, x=None, y=None):", src)
-        self.assertIn("return self._call(x=x, y=y)", src)
+        self.assertIn("return self._method_call(x=x, y=y)", src)
         self.assertEqual(len(list(method_to_call.named_modules())), 2)
         sess = self.check_ort(filename)
         input_names = [i.name for i in sess.get_inputs()]
@@ -197,7 +197,7 @@ def forward(self, x=None, y=None):
         self.assertExists(filename)
         src = method_to_call._method_src
         self.assertIn("f(self, x=None, y=None):", src)
-        self.assertIn("return self._call(x=x, y=y)", src)
+        self.assertIn("return self._method_call(x=x, y=y)", src)
         self.assertEqual(len(list(method_to_call.named_modules())), 2)
         sess = self.check_ort(filename)
         input_names = [i.name for i in sess.get_inputs()]
@@ -235,7 +235,7 @@ def forward(self, x, y=None):
         self.assertExists(filename)
         src = method_to_call._method_src
         self.assertIn("f(self, x, y=None):", src)
-        self.assertIn("return self._call(x=x, y=y)", src)
+        self.assertIn("return self._method_call(x=x, y=y)", src)
         self.assertEqual(len(list(method_to_call.named_modules())), 2)
         sess = self.check_ort(filename)
         input_names = [i.name for i in sess.get_inputs()]

_unittests/ut_export/test_dynamic_shapes.py

@@ -937,6 +937,31 @@ def test_invalid_dimensions_for_export(self):
             backed_size_oblivious = cpl.invalid_dimensions_for_export()
             self.assertFalse(backed_size_oblivious)
 
+    def test_guess_dynamic_shapes_missing(self):
+        class Model(torch.nn.Module):
+            def forward(self, x, y=None):
+                if y is None:
+                    return x.abs()
+                return x.abs() + y
+
+        model = Model()
+        x = torch.randn((5, 6))
+        y = model(x=x)
+        self.assertNotEmpty(y)
+
+        inputs = [
+            (tuple(), {"x": x}),
+            (tuple(), {"x": torch.randn((6, 6)), "y": torch.randn((6, 6))}),
+            (tuple(), {"x": torch.randn((7, 6)), "y": torch.randn((7, 6))}),
+        ]
+
+        mi = ModelInputs(model, inputs)
+        ds = mi.guess_dynamic_shapes()
+        DYN = torch.export.Dim.DYNAMIC
+        self.assertEqual(ds, ((), {"x": {0: DYN}, "y": {0: DYN}}))
+        _a, _kw, ds = mi.move_to_kwargs(*mi.inputs[-1], ds)
+        self.assertEqual(ds, (tuple(), {"x": {0: DYN}, "y": {0: DYN}}))
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/export/api.py

@@ -1,7 +1,7 @@
 import inspect
 import os
 import textwrap
-from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Sequence, Set, Tuple, Union
 import torch
 from .dynamic_shapes import ModelInputs
 from .onnx_plug import EagerDirectReplacementWithOnnx
@@ -340,6 +340,7 @@ def __init__(
         inline: bool = True,
         convert_after_n_calls: int = 2,
         patch_kwargs: Optional[Dict[str, Any]] = None,
+        skip_kwargs_names: Optional[Set[str]] = None,
     ):
         super().__init__()
         self._model_to_call = mod
@@ -354,6 +355,7 @@ def __init__(
         self._patch_kwargs = patch_kwargs
         self._method_src = None
         self.verbose = verbose
+        self.skip_kwargs_names = skip_kwargs_names
         self._to_onnx_kwargs = dict(
             input_names=input_names,
             target_opset=target_opset,
@@ -370,6 +372,7 @@ def __init__(
             onnx_plugs=onnx_plugs,
             inline=inline,
         )
+        self._export_done = False
 
     def __str__(self) -> str:
         return self.__repr__()
@@ -381,14 +384,28 @@ def __repr__(self) -> str:
         )
 
     def forward(self, *args, **kwargs):
-        self._inputs.append((args, kwargs))
-        if self.verbose:
-            print(
-                f"[method_to_onnx] input[{len(self._inputs)-1}]: "
-                f"{string_type((args, kwargs), with_shape=True)}"
+        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
+                        }
+                    ),
+                )
             )
-        if len(self._inputs) >= self._convert_after_n_calls:
-            self._convert_method_to_onnx()
+            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:
+                self._convert_method_to_onnx()
+                del self._inputs[:]
+                self._export_done = True
         return self._method_call(*args, **kwargs)
 
     def _convert_method_to_onnx(self):
@@ -473,6 +490,7 @@ def method_to_onnx(
     inline: bool = True,
     convert_after_n_calls: int = 2,
     patch_kwargs: Optional[Dict[str, Any]] = None,
+    skip_kwargs_names: Optional[Set[str]] = None,
 ) -> Callable:
     """
     Exports one method into ONNX for a module into ONNX.
@@ -499,8 +517,12 @@ def method_to_onnx(
     :param inline: inline local functions
     :param convert_after_n_calls: converts the model after this number of calls.
     :param patch_kwargs: patch arguments
+    :param skip_kwargs_names: use default values for these parameters part of
+        the signature of the method to export
     :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(
         mod=mod,
@@ -521,5 +543,6 @@ def method_to_onnx(
         inline=inline,
         convert_after_n_calls=convert_after_n_calls,
         patch_kwargs=patch_kwargs,
+        skip_kwargs_names=skip_kwargs_names,
     )
     return wrapped_model

onnx_diagnostic/export/dynamic_shapes.py

@@ -352,6 +352,19 @@ def _generic_walker_step(
                     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 @@ def _generic_walker_step(
             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,9 @@ def __init__(
             if self.signature
             else None
         )
+        self.forward_parameters_kinds = {
+            p.name: p.kind for p in self.signature.parameters.values()
+        }
         self.forward_ordered_parameter_names = (
             list(self.signature.parameters) if self.signature else None
         )
@@ -973,7 +991,13 @@ def guess_dynamic_shapes(self, auto: Union[bool, str] = False) -> DYNAMIC_SHAPES
             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,7 +1017,7 @@ def guess_dynamic_shapes(self, auto: Union[bool, str] = False) -> DYNAMIC_SHAPES
                 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",
@@ -1049,6 +1073,23 @@ def move_to_kwargs(
             _kw_dyn = kw_dyn
             kw_dyn = {}
             for name in self.forward_ordered_parameter_names:
+                if (
+                    self.forward_parameters_kinds[name] == inspect.Parameter.VAR_KEYWORD
+                    and name not in _kwargs
+                    and name in _kw_dyn
+                ):
+                    f = _kw_dyn[name]
+                    assert isinstance(
+                        f, dict
+                    ), f"Unexpected type for name={name!r}, _kw_dyn={_kw_dyn}"
+                    for _k, _w in f.items():
+                        assert (
+                            _k in _kwargs
+                        ), f"Parameter {_k!r} not in found in kwargs: {set(_kwargs)}"
+                        kwargs[_k] = _kwargs[_k]
+                        kw_dyn[_k] = f[_k]
+                    continue
+
                 if name in _kwargs:
                     kwargs[name] = _kwargs[name]
                 if name in _kw_dyn: