support more cases

Author: xadupre

_unittests/ut_export/test_dynamic_shapes.py

@@ -551,6 +551,57 @@ def test_couple_input_ds_cache(self):
                 ).invalid_paths(),
             )
 
+    def test_couple_input_ds_args_kwargs_0(self):
+        T3x1 = torch.rand((3, 1))
+        T3x4 = torch.rand((3, 4))
+        T5x6 = torch.rand((5, 6))
+        ds_batch = {0: "batch"}
+        ds_batch_seq = {0: "batch", 1: "seq"}
+        args = (T5x6,)
+        kwargs = {"A": T3x4, "B": (T3x1, T3x1)}
+        Cls = CoupleInputsDynamicShapes
+        self.assertEqual(
+            [], Cls(args, kwargs, {"A": ds_batch, "B": (ds_batch, ds_batch)}).invalid_paths()
+        )
+        self.assertEqual(
+            [],
+            Cls(
+                args, kwargs, {"A": ds_batch, "B": (ds_batch, ds_batch)}, args_names=["X"]
+            ).invalid_paths(),
+        )
+        self.assertEqual(
+            [("B", 1, "[1]")],
+            Cls(args, kwargs, {"A": ds_batch, "B": (ds_batch, ds_batch_seq)}).invalid_paths(),
+        )
+
+    def test_couple_input_ds_args_kwargs_1(self):
+        T3x1 = torch.rand((3, 1))
+        T3x4 = torch.rand((3, 4))
+        T5x1 = torch.rand((5, 1))
+        ds_batch = {0: "batch"}
+        ds_batch_seq = {0: "batch", 1: "seq"}
+        args = (T5x1,)
+        kwargs = {"A": T3x4, "B": (T3x1, T3x1)}
+        Cls = CoupleInputsDynamicShapes
+        self.assertEqual(
+            [],
+            Cls(
+                args,
+                kwargs,
+                {"X": ds_batch, "A": ds_batch, "B": (ds_batch, ds_batch)},
+                args_names=["X"],
+            ).invalid_paths(),
+        )
+        self.assertEqual(
+            [("X", "[1]"), ("B", 1, "[1]")],
+            Cls(
+                args,
+                kwargs,
+                {"X": ds_batch_seq, "A": ds_batch, "B": (ds_batch, ds_batch_seq)},
+                args_names=["X"],
+            ).invalid_paths(),
+        )
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/export/dynamic_shapes.py

@@ -455,14 +455,27 @@ def validate_inputs_for_export(
 class CoupleInputsDynamicShapes:
     """
     Pair inputs / dynamic shapes.
+
+    :param args: positional arguments
+    :param kwargs: named arguments
+    :param dynamic_shapes: dynamic shapes
+    :param args_names: if both args and kwargs are not empty, then
+        dynamic shapes must be a dictionary, and positional must be added
+        to the named arguments. Arguments names or a module must be given
+        in that case.
     """
 
     def __init__(
-        self, args: Tuple[Any, ...], kwargs: Dict[str, Any], dynamic_shapes: DYNAMIC_SHAPES
+        self,
+        args: Tuple[Any, ...],
+        kwargs: Dict[str, Any],
+        dynamic_shapes: DYNAMIC_SHAPES,
+        args_names: Optional[Union[torch.nn.Module, List[str]]] = None,
     ):
         self.args = args
         self.kwargs = kwargs
         self.dynamic_shapes = dynamic_shapes
+        self.args_names = args_names
 
     def __str__(self) -> str:
         return "\n".join(
@@ -497,7 +510,39 @@ def invalid_paths(self) -> List[Union[str, int]]:
                 f"dynamic_shapes={self.dynamic_shapes} should have the same type."
             )
             return list(self._valid_shapes(self.args, self.dynamic_shapes))
-        raise NotImplementedError("args and kwargs are filled, it is not implemented yet.")
+
+        assert isinstance(self.dynamic_shapes, dict), (
+            f"Both positional and named arguments (args and kwargs) are filled. "
+            f"dynamic shapes must a dictionary not {type(self.dynamic_shapes)}"
+        )
+        if not self.args_names and set(self.dynamic_shapes) & set(self.kwargs) == set(
+            self.dynamic_shapes
+        ):
+            # No dynamic shapes for the positional arguments.
+            return list(self._valid_shapes(self.kwargs, self.dynamic_shapes))
+
+        if isinstance(self.args_names, list):
+            if not set(self.args_names) & set(self.dynamic_shapes):
+                # No dynamic shapes for the positional arguments.
+                return list(self._valid_shapes(self.kwargs, self.dynamic_shapes))
+
+            assert self.args_names, (
+                "args and kwargs are filled, then args_names must be specified in "
+                "the constructor to move positional arguments to named arguments."
+            )
+            assert len(self.args) <= len(self.args_names), (
+                f"There are {len(self.args)} positional arguments "
+                f"but only {len(self.args_names)} names. "
+                f"args={string_type(self.args, with_shape=True)}, args_name={self.args_names}"
+            )
+            kwargs = dict(zip(self.args_names, self.args))
+            kwargs.update(self.kwargs)
+            return list(self._valid_shapes(kwargs, self.dynamic_shapes))
+
+        raise NotImplementedError(
+            f"Not yet implemented when args is filled, "
+            f"kwargs as well but args_names is {type(self.args_names)}"
+        )
 
     @classmethod
     def _valid_shapes(