refactor algorithm to validate shapes

Author: xadupre

_unittests/ut_export/test_dynamic_shapes.py

@@ -467,13 +467,13 @@ def test_couple_input_ds_0(self):
         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(({0: "d=[1]"},), Cls((T1x4,), {}, ({0: "batch"},)).invalid_paths())
+        self.assertEqual(({0: "d=[1]"},), Cls((T1x1,), {}, ({0: "batch"},)).invalid_paths())
         self.assertEqual(
-            [("A", "[0]")], Cls((), {"A": T1x1}, {"A": {0: "batch"}}).invalid_paths()
+            {"A": {0: "d=[1]"}}, Cls((), {"A": T1x1}, {"A": {0: "batch"}}).invalid_paths()
         )
         self.assertEqual(
-            [("A", "[0]")], Cls((), {"A": T1x4}, {"A": {0: "batch"}}).invalid_paths()
+            {"A": {0: "d=[1]"}}, Cls((), {"A": T1x4}, {"A": {0: "batch"}}).invalid_paths()
         )
 
     def test_couple_input_ds_1(self):
@@ -483,8 +483,10 @@ def test_couple_input_ds_1(self):
         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())
+        self.assertEqual(None, Cls(args, {}, (ds_batch, ds_batch)).invalid_paths())
+        self.assertEqual(
+            (None, {1: "d=[1]"}), Cls(args, {}, (ds_batch, ds_batch_seq)).invalid_paths()
+        )
 
     def test_couple_input_ds_2(self):
         T3x1 = torch.rand((3, 1))
@@ -493,9 +495,10 @@ def test_couple_input_ds_2(self):
         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(None, 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()
+            {"B": {1: "d=[1]"}},
+            Cls((), kwargs, {"A": ds_batch, "B": ds_batch_seq}).invalid_paths(),
         )
 
     def test_couple_input_ds_3(self):
@@ -506,10 +509,10 @@ def test_couple_input_ds_3(self):
         kwargs = {"A": T3x4, "B": (T3x1, T3x1)}
         Cls = CoupleInputsDynamicShapes
         self.assertEqual(
-            [], Cls((), kwargs, {"A": ds_batch, "B": (ds_batch, ds_batch)}).invalid_paths()
+            None, Cls((), kwargs, {"A": ds_batch, "B": (ds_batch, ds_batch)}).invalid_paths()
         )
         self.assertEqual(
-            [("B", 1, "[1]")],
+            {"B": (None, {1: "d=[1]"})},
             Cls((), kwargs, {"A": ds_batch, "B": (ds_batch, ds_batch_seq)}).invalid_paths(),
         )
 
@@ -532,15 +535,15 @@ def test_couple_input_ds_cache(self):
         Cls = CoupleInputsDynamicShapes
         with bypass_export_some_errors(patch_transformers=True):
             self.assertEqual(
-                [],
+                None,
                 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]")],
+                {"B": (None, [None, {2: "d=[1]"}, None, {2: "d=[1]"}])},
                 Cls(
                     (),
                     kwargs,
@@ -561,16 +564,16 @@ def test_couple_input_ds_args_kwargs_0(self):
         kwargs = {"A": T3x4, "B": (T3x1, T3x1)}
         Cls = CoupleInputsDynamicShapes
         self.assertEqual(
-            [], Cls(args, kwargs, {"A": ds_batch, "B": (ds_batch, ds_batch)}).invalid_paths()
+            None, Cls(args, kwargs, {"A": ds_batch, "B": (ds_batch, ds_batch)}).invalid_paths()
         )
         self.assertEqual(
-            [],
+            None,
             Cls(
                 args, kwargs, {"A": ds_batch, "B": (ds_batch, ds_batch)}, args_names=["X"]
             ).invalid_paths(),
         )
         self.assertEqual(
-            [("B", 1, "[1]")],
+            {"B": (None, {1: "d=[1]"})},
             Cls(args, kwargs, {"A": ds_batch, "B": (ds_batch, ds_batch_seq)}).invalid_paths(),
         )
 
@@ -584,7 +587,7 @@ def test_couple_input_ds_args_kwargs_1(self):
         kwargs = {"A": T3x4, "B": (T3x1, T3x1)}
         Cls = CoupleInputsDynamicShapes
         self.assertEqual(
-            [],
+            None,
             Cls(
                 args,
                 kwargs,
@@ -593,7 +596,7 @@ def test_couple_input_ds_args_kwargs_1(self):
             ).invalid_paths(),
         )
         self.assertEqual(
-            [("X", "[1]"), ("B", 1, "[1]")],
+            {"X": {1: "d=[1]"}, "B": (None, {1: "d=[1]"})},
             Cls(
                 args,
                 kwargs,

onnx_diagnostic/export/dynamic_shapes.py

@@ -488,7 +488,7 @@ def __str__(self) -> str:
             ]
         )
 
-    def invalid_paths(self) -> List[Union[str, int]]:
+    def invalid_paths(self) -> Any:
         """
         Tells the inputs are valid based on the dynamic shapes definition.
         The method assumes that all custom classes can be serialized.
@@ -498,18 +498,42 @@ def invalid_paths(self) -> List[Union[str, int]]:
         The function checks that a dynamic dimension does not receive a value
         of 0 or 1. It returns a list of invalid path.
         """
+        return self._generic_walker(self._valid_shapes_tensor)
+
+    @classmethod
+    def _valid_shapes_tensor(cls, inputs: Any, ds: Any) -> Iterable:
+        assert isinstance(inputs, torch.Tensor), f"unexpected type for inputs {type(inputs)}"
+        assert isinstance(ds, dict) and all(isinstance(s, int) for s in ds), (
+            f"Unexpected types, inputs is a Tensor but ds is {ds}, "
+            f"a dictionary is expected to specify a dimension dimension"
+        )
+        issues = {}
+        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
+                    issues[i] = f"d=[{d}]"
+        return issues if issues else None
+
+    def _generic_walker(self, method_to_call: Callable) -> Any:
+        """
+        Generic deserializator walking through inputs and dynamic_shapes all along.
+        The function returns a result with the same structure as the dynamic shapes.
+        """
         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))
+            return self._generic_walker_step(method_to_call, 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))
+            return self._generic_walker_step(method_to_call, self.args, self.dynamic_shapes)
 
         assert isinstance(self.dynamic_shapes, dict), (
             f"Both positional and named arguments (args and kwargs) are filled. "
@@ -519,12 +543,14 @@ def invalid_paths(self) -> List[Union[str, int]]:
             self.dynamic_shapes
         ):
             # No dynamic shapes for the positional arguments.
-            return list(self._valid_shapes(self.kwargs, self.dynamic_shapes))
+            return self._generic_walker_step(method_to_call, 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))
+                return self._generic_walker_step(
+                    method_to_call, self.kwargs, self.dynamic_shapes
+                )
 
             assert self.args_names, (
                 "args and kwargs are filled, then args_names must be specified in "
@@ -537,62 +563,52 @@ def invalid_paths(self) -> List[Union[str, int]]:
             )
             kwargs = dict(zip(self.args_names, self.args))
             kwargs.update(self.kwargs)
-            return list(self._valid_shapes(kwargs, self.dynamic_shapes))
+            return self._generic_walker_step(method_to_call, 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(
-        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}"
+    def _generic_walker_step(cls, method_to_call: Callable, inputs: Any, ds: Any) -> Iterable:
         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.
-                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."
+            return method_to_call(inputs, ds)
+        if isinstance(inputs, (int, float, str)):
+            return None
+        if isinstance(inputs, (tuple, list, dict)):
+            assert type(ds) is type(
+                inputs
+            ), f"Type mismatch between inputs {type(inputs)} and ds={type(ds)}"
+            assert len(ds) == len(inputs), (
+                f"Length mismatch between inputs {len(inputs)} "
+                f"and ds={len(ds)}\n"
+                f"inputs={string_type(inputs, with_shape=True)}, ds={ds}"
+            )
+            if isinstance(inputs, (tuple, list)):
+                value = []
+                for i, d in zip(inputs, ds):
+                    value.append(cls._generic_walker_step(method_to_call, i, d))
+                return (
+                    (value if isinstance(ds, list) else tuple(value))
+                    if any(v is not None for v in value)
+                    else None
                 )
-                flat, _spec = torch.utils._pytree.tree_flatten(inputs)
-                for path in cls._valid_shapes(
-                    flat, ds, prefix=(*prefix, inputs.__class__.__name__)
-                ):
-                    yield path
+            assert set(inputs) == set(
+                ds
+            ), f"Keys mismatch between inputs {set(inputs)} and ds={set(ds)}"
+            dvalue = {}
+            for k, v in inputs.items():
+                t = cls._generic_walker_step(method_to_call, v, ds[k])
+                if t is not None:
+                    dvalue[k] = t
+            return dvalue if dvalue else None
+
+        # 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)
+        return cls._generic_walker_step(method_to_call, flat, ds)