merge

Author: xadupre

_unittests/ut_export/test_dynamic_shapes.py

@@ -5,6 +5,10 @@
 from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache
 from onnx_diagnostic.export import ModelInputs
 from onnx_diagnostic.export.dynamic_shapes import CoupleInputsDynamicShapes
+<<<<<<< HEAD
+=======
+from onnx_diagnostic.torch_export_patches import bypass_export_some_errors
+>>>>>>> abbcc6bca05bb6d921592fe814b92567b36da435
 
 
 class TestDynamicShapes(ExtTestCase):
@@ -529,23 +533,75 @@ def test_couple_input_ds_cache(self):
 
         kwargs = {"A": T3x4, "B": (T3x1, cache)}
         Cls = CoupleInputsDynamicShapes
+        with bypass_export_some_errors(patch_transformers=True):
+            self.assertEqual(
+                [],
+                Cls(
+                    (),
+                    kwargs,
+                    {"A": ds_batch, "B": (ds_batch, [ds_batch, ds_batch, ds_batch, ds_batch])},
+                ).invalid_paths(),
+            )
+            self.assertEqual(
+                [("B", 1, "DynamicCache", 1, "[2]"), ("B", 1, "DynamicCache", 3, "[2]")],
+                Cls(
+                    (),
+                    kwargs,
+                    {
+                        "A": ds_batch,
+                        "B": (ds_batch, [ds_batch, ds_batch_seq, ds_batch, ds_batch_seq]),
+                    },
+                ).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,
-                {"A": ds_batch, "B": (ds_batch, [ds_batch, ds_batch, ds_batch, ds_batch])},
+                {"X": ds_batch, "A": ds_batch, "B": (ds_batch, ds_batch)},
+                args_names=["X"],
             ).invalid_paths(),
         )
         self.assertEqual(
-            [("B", 1, "DynamicCache", 1, "[2]"), ("B", 1, "DynamicCache", 3, "[2]")],
+            [("X", "[1]"), ("B", 1, "[1]")],
             Cls(
-                (),
+                args,
                 kwargs,
-                {
-                    "A": ds_batch,
-                    "B": (ds_batch, [ds_batch, ds_batch_seq, ds_batch, ds_batch_seq]),
-                },
+                {"X": ds_batch_seq, "A": ds_batch, "B": (ds_batch, ds_batch_seq)},
+                args_names=["X"],
             ).invalid_paths(),
         )
 

_unittests/ut_torch_onnx/test_sbs.py

@@ -18,7 +18,11 @@ class TestSideBySide(ExtTestCase):
 
     @hide_stdout()
     @unittest.skipIf(to_onnx is None, "to_onnx not installed")
+<<<<<<< HEAD
     @ignore_errors(OSError, "connectivity issues")
+=======
+    @ignore_errors(OSError)  # connectivity issues
+>>>>>>> abbcc6bca05bb6d921592fe814b92567b36da435
     @ignore_warnings((UserWarning,))
     def test_ep_onnx_sync_exp(self):
         import torch
@@ -52,7 +56,6 @@ def forward(self, x):
 
     @hide_stdout()
     @ignore_warnings((DeprecationWarning, FutureWarning, UserWarning))
-    @unittest.skipIf(to_onnx is None, "to_onnx not installed")
     def test_ep_onnx_sync_a(self):
         import torch
 
@@ -69,9 +72,10 @@ def forward(self, x):
         ep = torch.export.export(
             Model(), (x,), dynamic_shapes=({0: torch.export.Dim("batch")},)
         )
-        onx = torch.onnx.export(
-            Model(), (x,), dynamic_shapes=({0: torch.export.Dim("batch")},), dynamo=True
-        ).model_proto
+        epo = torch.onnx.export(
+            ep, (x,), dynamic_shapes=({0: torch.export.Dim("batch")},), dynamo=True
+        )
+        onx = epo.model_proto
         results = list(
             run_aligned(
                 ep,

onnx_diagnostic/export/dynamic_shapes.py

@@ -437,7 +437,7 @@ def move_to_kwargs(
 
     def validate_inputs_for_export(
         self, dynamic_shapes: Optional[DYNAMIC_SHAPES] = None
-    ) -> List[List[str]]:
+    ) -> List[List[Union[int, str]]]:
         """
         Validates the inputs the class contains for the given dynamic shapes.
         If not specified, the dynamic_shapes are guessed.
@@ -447,22 +447,35 @@ def validate_inputs_for_export(
         """
         if dynamic_shapes is None:
             if len(self.inputs) == 1:
-                return True
+                return []
             dyn_shapes = self.guess_dynamic_shapes()
         return [CoupleInputsDynamicShapes(*i, dyn_shapes).invalid_paths() for i in self.inputs]
 
 
 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(
@@ -541,6 +586,11 @@ def _valid_shapes(
                             yield path
             else:
                 # A custom class.
+                assert 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."
+                )
                 flat, _spec = torch.utils._pytree.tree_flatten(inputs)
                 for path in cls._valid_shapes(
                     flat, ds, prefix=(*prefix, inputs.__class__.__name__)

onnx_diagnostic/torch_models/test_helper.py

@@ -1,3 +1,4 @@
+import datetime
 import inspect
 import os
 from typing import Any, Dict, List, Optional, Tuple, Union
@@ -139,6 +140,55 @@ def _make_folder_name(
     return "-".join(els)
 
 
+def version_summary() -> Dict[str, Union[int, float, str]]:
+    """
+    Example:
+
+    .. runpython::
+        :showcode:
+
+        import pprint
+        from onnx_diagnostic.torch_models.test_helper import version_summary
+
+        pprint.pprint(version_summary())
+    """
+    import numpy
+
+    summary: Dict[str, Union[int, float, str]] = {
+        "version_torch": torch.__version__,
+        "version_numpy": numpy.__version__,
+    }
+    try:
+        import transformers
+
+        summary["version_transformers"] = transformers.__version__
+    except ImportError:
+        pass
+    try:
+        import onnx
+
+        summary["version_onnx"] = onnx.__version__
+    except ImportError:
+        pass
+    try:
+        import onnxscript
+
+        summary["version_onnxscript"] = onnxscript.__version__
+    except ImportError:
+        pass
+    try:
+        import onnxruntime
+
+        summary["version_onnxruntime"] = onnxruntime.__version__
+    except ImportError:
+        pass
+    import onnx_diagnostic
+
+    summary["version_onnx_diagnostic"] = onnx_diagnostic.__version__
+    summary["version_date"] = datetime.datetime.now().strftime("%Y-%m-%dT%H:%M:%S")
+    return summary
+
+
 def validate_model(
     model_id: str,
     task: Optional[str] = None,
@@ -180,7 +230,7 @@ def validate_model(
         another one with whatever the function produces
     """
     assert not trained, f"trained={trained} not supported yet"
-    summary: Dict[str, Union[int, float, str]] = {}
+    summary = version_summary()
     if dump_folder:
         folder_name = _make_folder_name(
             model_id, exporter, optimization, dtype=dtype, device=device