mypy

Author: xadupre

CHANGELOGS.rst

@@ -4,6 +4,7 @@ Change Logs
 0.8.8
 +++++
 
+* :pr:`375`: export a method to onnx in order to export using method generate
 * :pr:`372`: fix patch on rotary embedding
 * :pr:`371`: fix make_fake_with_dynamic_dimensions
 

onnx_diagnostic/export/api.py

@@ -347,7 +347,7 @@ def __init__(
         self._call = (
             self._model_to_call if method_name == "forward" else getattr(mod, method_name)
         )
-        self._inputs = []
+        self._inputs: List[Tuple[Tuple[Any, ...], Dict[str, Any]]] = []
         self._convert_after_n_calls = convert_after_n_calls
         self._patch_kwargs = patch_kwargs
         self._method_src = None

onnx_diagnostic/export/dynamic_shapes.py

@@ -1000,7 +1000,10 @@ def guess_dynamic_shapes(self, auto: Union[bool, str] = False) -> DYNAMIC_SHAPES
                 msg=lambda name=name: f" failing input {name!r}",
             )
         # reordering
-        kwargs = {p: kwargs[p] for p in self.forward_ordered_parameter_names if p in kwargs}
+        if kwargs is not None:
+            kwargs = {
+                p: kwargs[p] for p in self.forward_ordered_parameter_names if p in kwargs
+            }
         return tuple(args), kwargs
 
     def move_to_kwargs(
@@ -1063,8 +1066,14 @@ def move_to_kwargs(
                 f"and kwargs={set(kwargs)}, "
                 f"forward_ordered_parameter_names={self.forward_ordered_parameter_names}"
             )
-        kwargs = {p: kwargs[p] for p in self.forward_ordered_parameter_names if p in kwargs}
-        kw_dyn = {p: kw_dyn[p] for p in self.forward_ordered_parameter_names if p in kw_dyn}
+        if kwargs is not None:
+            kwargs = {
+                p: kwargs[p] for p in self.forward_ordered_parameter_names if p in kwargs
+            }
+        if kw_dyn is not None:
+            kw_dyn = {
+                p: kw_dyn[p] for p in self.forward_ordered_parameter_names if p in kw_dyn
+            }
         return args, kwargs, (tuple(), kw_dyn)
 
     def validate_inputs_for_export(