valid inputs

Author: xadupre

_unittests/ut_export/test_dynamic_shapes.py

@@ -4,6 +4,7 @@
 from onnx_diagnostic.helpers import string_type
 from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache
 from onnx_diagnostic.export import ModelInputs
+from onnx_diagnostic.export.dynamic_shapes import CoupleInputsDynamicShapes
 
 
 class TestDynamicShapes(ExtTestCase):
@@ -453,6 +454,101 @@ def forward(self, cache, z):
             ),
         )
 
+    def test_couple_input_ds_0(self):
+        T3x4 = torch.rand((3, 4))
+        T3x1 = torch.rand((3, 1))
+        Cls = CoupleInputsDynamicShapes
+        self.assertEmpty(Cls((T3x4,), {}, ({0: "batch"},)).invalid_paths())
+        self.assertEmpty(Cls((T3x1,), {}, ({0: "batch"},)).invalid_paths())
+        self.assertEmpty(Cls((), {"A": T3x1}, {"A": {0: "batch"}}).invalid_paths())
+        self.assertEmpty(Cls((), {"A": T3x4}, {"A": {0: "batch"}}).invalid_paths())
+
+        T1x4 = torch.rand((1, 4))
+        T1x1 = torch.rand((1, 1))
+        Cls = CoupleInputsDynamicShapes
+        self.assertEqual([(0, "[0]")], Cls((T1x4,), {}, ({0: "batch"},)).invalid_paths())
+        self.assertEqual([(0, "[0]")], Cls((T1x1,), {}, ({0: "batch"},)).invalid_paths())
+        self.assertEqual(
+            [("A", "[0]")], Cls((), {"A": T1x1}, {"A": {0: "batch"}}).invalid_paths()
+        )
+        self.assertEqual(
+            [("A", "[0]")], Cls((), {"A": T1x4}, {"A": {0: "batch"}}).invalid_paths()
+        )
+
+    def test_couple_input_ds_1(self):
+        T3x1 = torch.rand((3, 1))
+        T3x4 = torch.rand((3, 4))
+        ds_batch = {0: "batch"}
+        ds_batch_seq = {0: "batch", 1: "seq"}
+        args = (T3x4, T3x1)
+        Cls = CoupleInputsDynamicShapes
+        self.assertEqual([], Cls(args, {}, (ds_batch, ds_batch)).invalid_paths())
+        self.assertEqual([(1, "[1]")], Cls(args, {}, (ds_batch, ds_batch_seq)).invalid_paths())
+
+    def test_couple_input_ds_2(self):
+        T3x1 = torch.rand((3, 1))
+        T3x4 = torch.rand((3, 4))
+        ds_batch = {0: "batch"}
+        ds_batch_seq = {0: "batch", 1: "seq"}
+        kwargs = {"A": T3x4, "B": T3x1}
+        Cls = CoupleInputsDynamicShapes
+        self.assertEqual([], Cls((), kwargs, {"A": ds_batch, "B": ds_batch}).invalid_paths())
+        self.assertEqual(
+            [("B", "[1]")], Cls((), kwargs, {"A": ds_batch, "B": ds_batch_seq}).invalid_paths()
+        )
+
+    def test_couple_input_ds_3(self):
+        T3x1 = torch.rand((3, 1))
+        T3x4 = torch.rand((3, 4))
+        ds_batch = {0: "batch"}
+        ds_batch_seq = {0: "batch", 1: "seq"}
+        kwargs = {"A": T3x4, "B": (T3x1, T3x1)}
+        Cls = CoupleInputsDynamicShapes
+        self.assertEqual(
+            [], Cls((), kwargs, {"A": ds_batch, "B": (ds_batch, ds_batch)}).invalid_paths()
+        )
+        self.assertEqual(
+            [("B", 1, "[1]")],
+            Cls((), kwargs, {"A": ds_batch, "B": (ds_batch, ds_batch_seq)}).invalid_paths(),
+        )
+
+    def test_couple_input_ds_cache(self):
+        T3x1 = torch.rand((3, 1))
+        T3x4 = torch.rand((3, 4))
+        ds_batch = {0: "batch"}
+        ds_batch_seq = {0: "batch", 2: "seq"}
+
+        n_layers = 2
+        bsize, nheads, slen, dim = 2, 4, 1, 7
+        cache = make_dynamic_cache(
+            [
+                (torch.randn(bsize, nheads, slen, dim), torch.randn(bsize, nheads, slen, dim))
+                for i in range(n_layers)
+            ]
+        )
+
+        kwargs = {"A": T3x4, "B": (T3x1, cache)}
+        Cls = CoupleInputsDynamicShapes
+        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(),
+        )
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

_unittests/ut_torch_onnx/test_sbs.py

@@ -1,5 +1,10 @@
 import unittest
-from onnx_diagnostic.ext_test_case import ExtTestCase, hide_stdout, ignore_warnings
+from onnx_diagnostic.ext_test_case import (
+    ExtTestCase,
+    hide_stdout,
+    ignore_warnings,
+    ignore_errors,
+)
 from onnx_diagnostic.reference import ExtendedReferenceEvaluator
 from onnx_diagnostic.torch_onnx.sbs import run_aligned
 
@@ -13,6 +18,7 @@ class TestSideBySide(ExtTestCase):
 
     @hide_stdout()
     @unittest.skipIf(to_onnx is None, "to_onnx not installed")
+    @ignore_errors(OSError, "connectivity issues")
     @ignore_warnings((UserWarning,))
     def test_ep_onnx_sync_exp(self):
         import torch

onnx_diagnostic/export/dynamic_shapes.py

@@ -1,9 +1,11 @@
 import inspect
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union
 import numpy as np
 import torch
 from ..helpers import string_type
 
+DYNAMIC_SHAPES = Tuple[Tuple[Any, ...], Dict[str, Any]]
+
 
 class ModelInputs:
     """
@@ -218,7 +220,7 @@ def process_inputs(
         return new_inputs
 
     @property
-    def true_model_name(self):
+    def true_model_name(self) -> str:
         "Returns class name or module name."
         return (
             self.model.__class__.__name__
@@ -227,7 +229,7 @@ def true_model_name(self):
         )
 
     @property
-    def full_name(self):
+    def full_name(self) -> str:
         "Returns a name and class name."
         if self.method_name == "forward":
             return f"{self.name}:{self.true_model_name}"
@@ -337,9 +339,7 @@ def guess_dynamic_shape_object(self, *objs: Any, msg: Optional[Callable] = None)
             f"{string_type(objs)}{msg() if msg else ''} in {self.module_name_type}"
         )
 
-    def guess_dynamic_shapes(
-        self,
-    ) -> Tuple[Tuple[Any, ...], Dict[str, Any]]:
+    def guess_dynamic_shapes(self) -> DYNAMIC_SHAPES:
         """
         Guesses the dynamic shapes for that module from two execution.
         If there is only one execution, then that would be static dimensions.
@@ -386,7 +386,7 @@ def move_to_kwargs(
         args: Tuple[Any, ...],
         kwargs: Dict[str, Any],
         dynamic_shapes: Tuple[Tuple[Any, ...], Dict[str, Any]],
-    ) -> Tuple[Tuple[Any, ...], Dict[str, Any], Tuple[Tuple[Any, ...], Dict[str, Any]]]:
+    ) -> Tuple[Tuple[Any, ...], Dict[str, Any], DYNAMIC_SHAPES]:
         """
         Uses the signatures to move positional arguments (args) to named arguments (kwargs)
         with the corresponding dynamic shapes.
@@ -434,3 +434,115 @@ def move_to_kwargs(
                 f"forward_ordered_parameter_names={self.forward_ordered_parameter_names}"
             )
         return args, kwargs, (tuple(), kw_dyn)
+
+    def validate_inputs_for_export(
+        self, dynamic_shapes: Optional[DYNAMIC_SHAPES] = None
+    ) -> List[List[str]]:
+        """
+        Validates the inputs the class contains for the given dynamic shapes.
+        If not specified, the dynamic_shapes are guessed.
+
+        :param dynamic_shapes: dynamic shapes to validate
+        :return: a list of lists, every list contains the path the invalid dimension
+        """
+        if dynamic_shapes is None:
+            if len(self.inputs) == 1:
+                return True
+            dyn_shapes = self.guess_dynamic_shapes()
+        return [CoupleInputsDynamicShapes(*i, dyn_shapes).invalid_paths() for i in self.inputs]
+
+
+class CoupleInputsDynamicShapes:
+    """
+    Pair inputs / dynamic shapes.
+    """
+
+    def __init__(
+        self, args: Tuple[Any, ...], kwargs: Dict[str, Any], dynamic_shapes: DYNAMIC_SHAPES
+    ):
+        self.args = args
+        self.kwargs = kwargs
+        self.dynamic_shapes = dynamic_shapes
+
+    def __str__(self) -> str:
+        return "\n".join(
+            [
+                f"{self.__class__.__name__}(",
+                f"    args={string_type(self.args, with_shape=True)},"
+                f"    kwargs={string_type(self.kwargs, with_shape=True)},"
+                f"    dynamic_shapes={string_type(self.dynamic_shapes, with_shape=True)},"
+                f")",
+            ]
+        )
+
+    def invalid_paths(self) -> List[Union[str, int]]:
+        """
+        Tells the inputs are valid based on the dynamic shapes definition.
+        The method assumes that all custom classes can be serialized.
+        If some patches were applied to export, they should enabled while
+        calling this method if the inputs contains such classes.
+
+        The function checks that a dynamic dimension does not receive a value
+        of 0 or 1. It returns a list of invalid path.
+        """
+        if not self.args:
+            assert isinstance(self.kwargs, dict) and isinstance(self.dynamic_shapes, dict), (
+                f"Type mismatch, args={string_type(self.args)} and "
+                f"dynamic_shapes={self.dynamic_shapes} should have the same type."
+            )
+            return list(self._valid_shapes(self.kwargs, self.dynamic_shapes))
+        if not self.kwargs:
+            assert isinstance(self.args, tuple) and isinstance(self.dynamic_shapes, tuple), (
+                f"Type mismatch, args={string_type(self.args)} and "
+                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.")
+
+    @classmethod
+    def _valid_shapes(
+        cls, inputs: Any, ds: Any, prefix: Tuple[Union[int, str], ...] = ()
+    ) -> Iterable:
+        assert all(isinstance(i, (int, str)) for i in prefix), f"Unexpected prefix {prefix}"
+        if isinstance(inputs, torch.Tensor):
+            assert isinstance(ds, dict) and all(
+                isinstance(s, int) for s in ds
+            ), f"Unexpected types, inputs is a Tensor but ds={ds}, prefix={prefix}"
+            for i, d in enumerate(inputs.shape):
+                if i in ds and not isinstance(ds[i], int):
+                    # dynamic then
+                    if d in {0, 1}:
+                        # export issues for sure
+                        yield (*prefix, f"[{i}]")
+        else:
+            if isinstance(inputs, (int, float, str)):
+                pass
+            elif isinstance(inputs, (tuple, list, dict)):
+                assert type(ds) is type(inputs), (
+                    f"Type mismatch between inputs {type(inputs)} "
+                    f"and ds={type(ds)}, prefix={prefix!r}"
+                )
+                assert len(ds) == len(inputs), (
+                    f"Length mismatch between inputs {len(inputs)} "
+                    f"and ds={len(ds)}, prefix={prefix!r}\n"
+                    f"inputs={string_type(inputs, with_shape=True)}, ds={ds}"
+                )
+                if isinstance(inputs, (tuple, list)):
+                    for ind, (i, d) in enumerate(zip(inputs, ds)):
+                        for path in cls._valid_shapes(i, d, prefix=(*prefix, ind)):
+                            yield path
+                else:
+                    assert set(inputs) == set(ds), (
+                        f"Keys mismatch between inputs {set(inputs)} "
+                        f"and ds={set(ds)}, prefix={prefix!r}"
+                    )
+                    for k, v in inputs.items():
+                        for path in cls._valid_shapes(v, ds[k], prefix=(*prefix, k)):
+                            yield path
+            else:
+                # A custom class.
+                flat, _spec = torch.utils._pytree.tree_flatten(inputs)
+                for path in cls._valid_shapes(
+                    flat, ds, prefix=(*prefix, inputs.__class__.__name__)
+                ):
+                    yield path