onnx-diagnostic 0.2.1__py3-none-any.whl → 0.3.0__py3-none-any.whl

onnx_diagnostic/export/dynamic_shapes.py

@@ -3,6 +3,344 @@ from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 import numpy as np
 import torch
 from ..helpers import string_type
+from ..helpers.cache_helper import flatten_unflatten_for_dynamic_shapes
+
+DYNAMIC_SHAPES = Tuple[Tuple[Any, ...], Dict[str, Any]]
+
+
+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,
+        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(
+            [
+                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 replace_string_by(self, value: Any = None):
+        """
+        Replaces string by the value ``torch.export.Dim.DYNAMIC``
+        (default) or any other value specified by value.
+
+        Example:
+
+        .. runpython::
+            :showcode:
+
+            import torch
+            from onnx_diagnostic.export.dynamic_shapes import CoupleInputsDynamicShapes
+
+            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)}
+            ds = {"A": ds_batch, "B": (ds_batch, ds_batch_seq)}
+            print(CoupleInputsDynamicShapes((), kwargs, ds).replace_string_by())
+        """
+        return self._generic_walker(
+            lambda inputs, ds, value=value: self._replace_string_dim_tensor(
+                inputs, ds, value=value
+            )
+        )
+
+    @classmethod
+    def _replace_string_dim_tensor(cls, inputs, ds, value=None):
+        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"
+        )
+        if value is None:
+            value = torch.export.Dim.DYNAMIC
+        new_ds = ds.copy()
+        for i, v in ds.items():
+            if isinstance(v, str):
+                new_ds[i] = value
+        return new_ds
+
+    def invalid_dimensions_for_export(self):
+        """
+        Tells if 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 the unexpected values in the same structure as
+        the given dynamic shapes.
+
+        Example:
+
+        .. runpython::
+            :showcode:
+
+            import torch
+            from onnx_diagnostic.export.dynamic_shapes import CoupleInputsDynamicShapes
+
+            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)}
+            ds = {"A": ds_batch, "B": (ds_batch, ds_batch_seq)}
+            print(CoupleInputsDynamicShapes((), kwargs, ds).invalid_dimensions_for_export())
+
+        In case it works, it shows:
+
+        .. runpython::
+            :showcode:
+
+            import torch
+            from onnx_diagnostic.export.dynamic_shapes import CoupleInputsDynamicShapes
+
+            T3x2 = torch.rand((3, 2))
+            T3x4 = torch.rand((3, 4))
+            ds_batch = {0: "batch"}
+            ds_batch_seq = {0: "batch", 1: "seq"}
+            kwargs = {"A": T3x4, "B": (T3x2, T3x2)}
+            ds = {"A": ds_batch, "B": (ds_batch, ds_batch_seq)}
+            print(CoupleInputsDynamicShapes((), kwargs, ds).invalid_dimensions_for_export())
+        """
+        return self._generic_walker(self._valid_shapes_tensor)
+
+    @classmethod
+    def _valid_shapes_tensor(cls, inputs, ds):
+        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, processor: Callable):
+        """
+        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 self._generic_walker_step(processor, 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 self._generic_walker_step(processor, self.args, self.dynamic_shapes)
+
+        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 self._generic_walker_step(processor, 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 self._generic_walker_step(processor, 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 self._generic_walker_step(processor, 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 _generic_walker_step(cls, processor: Callable, inputs, ds):
+        if isinstance(inputs, torch.Tensor):
+            return processor(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(processor, i, d))
+                return (
+                    (value if isinstance(ds, list) else tuple(value))
+                    if any(v is not None for v in value)
+                    else None
+                )
+            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(processor, 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(processor, flat, ds)
+
+    class ChangeDimensionProcessor:
+        def __init__(self):
+            self.mapping = {}
+
+        def _build_new_shape(
+            self, shape: Tuple[int, ...], ds: Dict[int, Any]
+        ) -> Tuple[int, ...]:
+            new_shape = list(shape)
+            for i in range(len(shape)):
+                if i in ds:
+                    if isinstance(ds[i], str):
+                        d = ds[i]
+                    elif isinstance(
+                        ds[i],
+                        (
+                            torch.export.dynamic_shapes._DerivedDim,
+                            torch.export.dynamic_shapes._Dim,
+                        ),
+                    ):
+                        d = str(ds[i])
+                    elif not isinstance(ds[i], int):
+                        raise NotImplementedError(f"Unable to handle type {ds[i]} in {ds}")
+                    if d in self.mapping:
+                        new_dim = self.mapping[d]
+                    else:
+                        new_dim = shape[i] + 1
+                        self.mapping[d] = new_dim
+                    new_shape[i] = new_dim
+            return tuple(new_shape)
+
+        def _build_new_tensor(self, tensor: torch.Tensor, new_shape: Tuple[int, ...]):
+            rank = len(tensor.shape)
+            for i in range(len(tensor.shape)):
+                d0 = tensor.shape[i]
+                d1 = new_shape[i]
+                if d0 == d1:
+                    continue
+                alt_shape = list(tensor.shape)
+                alt_shape[i] = d1
+                new_tensor = torch.zeros(
+                    tuple(alt_shape), dtype=tensor.dtype, device=tensor.device
+                )
+                mind = min(d0, d1)
+                indices = [slice(None) for _ in range(rank)]
+                indices[i] = slice(0, mind)
+                ind = tuple(indices)
+                new_tensor[ind] = tensor[ind]
+                if d1 > mind:
+                    for k in range(d1 - mind):
+                        indices0 = [slice(None) for _ in range(rank)]
+                        indices1 = [slice(None) for _ in range(rank)]
+                        indices1[i] = mind + k
+                        indices0[i] = k % mind
+                        new_tensor[tuple(indices1)] = tensor[tuple(indices0)]
+                tensor = new_tensor
+            return tensor
+
+        def __call__(self, inputs, ds):
+            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"
+            )
+            new_shape = self._build_new_shape(inputs.shape, ds)
+            return self._build_new_tensor(inputs, new_shape)
+
+    def change_dynamic_dimensions(self):
+        """
+        A model exported with dynamic shapes is not necessarily dynamic
+        just because the user specified dynamic shapes. The algorithm
+        may discover that a dimension cannot be dynamic and then continues
+        the export making the assumption it is static. That may lead a wrong
+        model. This function produces a new set of inputs with different values
+        for the dimension than the first ones, assuming they were used to export
+        the model.
+
+        Example:
+
+        .. runpython::
+            :showcode:
+
+            import torch
+            from onnx_diagnostic.helpers import string_type
+            from onnx_diagnostic.export.dynamic_shapes import CoupleInputsDynamicShapes
+
+            T3x15 = torch.rand((3, 15))
+            T3x20 = torch.rand((3, 20))
+            T3x4 = torch.rand((3, 4))
+            ds_batch = {0: "batch"}
+            ds_batch_seq = {0: "batch", 1: "seq"}
+            kwargs = {"A": T3x4, "B": (T3x15, T3x20)}
+            ds = {"A": ds_batch, "B": (ds_batch, ds_batch_seq)}
+            new_kwargs = CoupleInputsDynamicShapes((), kwargs, ds).change_dynamic_dimensions()
+            print("before:", string_type(kwargs, with_shape=True))
+            print("-after:", string_type(new_kwargs, with_shape=True))
+        """
+        return self._generic_walker(self.ChangeDimensionProcessor())
 
 
 class ModelInputs:
@@ -69,7 +407,7 @@ class ModelInputs:
         ds = mi.guess_dynamic_shapes()
         pprint.pprint(ds)
 
-    **and and kwargs**
+    **args and kwargs**
 
     .. runpython::
         :showcode:
@@ -218,7 +556,7 @@ class ModelInputs:
         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 +565,7 @@ class ModelInputs:
         )
 
     @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}"
@@ -240,10 +578,27 @@ class ModelInputs:
             return f"type({self.name})={self.true_model_name}"
         return f"type({self.name})={self.true_model_name}.{self.method_name}"
 
-    def guess_dynamic_dimensions(self, *tensors) -> Dict[int, Any]:
-        """Infers the dynamic dimension from multiple shapes."""
+    def guess_dynamic_dimensions(
+        self, *tensors, auto: bool = False
+    ) -> Optional[Dict[int, Any]]:
+        """
+        Infers the dynamic dimension from multiple shapes.
+        If auto is True, it returns ``torch.export.Dim.AUTO`` for every dimension
+        which cannot be guessed. Two tensors with the same value for one dimension
+        can be guessed, but if there is only 1, it cannot.
+        """
         if len(tensors) == 1:
-            return {}
+            if isinstance(tensors[0], (int, float)):
+                return None
+            assert isinstance(tensors[0], torch.Tensor), (
+                f"Unexpected type for tensors {string_type(tensors, with_shape=True)}, "
+                f"Only tensors are allowed."
+            )
+            return (
+                {i: torch.export.Dim.AUTO for i in range(len(tensors[0].shape))}  # noqa: C420
+                if auto
+                else {}
+            )
         shapes = [t.shape for t in tensors]
         set_length = set(len(s) for s in shapes)
         assert len(set_length) == 1, (
@@ -265,7 +620,9 @@ class ModelInputs:
                 continue
         return res
 
-    def guess_dynamic_shape_object(self, *objs: Any, msg: Optional[Callable] = None) -> Any:
+    def guess_dynamic_shape_object(
+        self, *objs: Any, auto: bool = False, msg: Optional[Callable] = None
+    ) -> Any:
         """Guesses the dynamic shapes for one argument."""
         if len(objs) == 0:
             return None
@@ -279,7 +636,7 @@ class ModelInputs:
         if isinstance(obj, (bool, int, float, str)):
             return None
         if isinstance(obj, (torch.Tensor, np.ndarray)):
-            return self.guess_dynamic_dimensions(*objs)
+            return self.guess_dynamic_dimensions(*objs, auto=auto)
 
         if isinstance(obj, tuple):
             kl = set(len(o) for o in objs)
@@ -288,7 +645,9 @@ class ModelInputs:
             ), f"Unexpected variety of tuple lengths {kl}{msg() if msg else ''}"
             shapes: Any = []
             for i in range(kl.pop()):
-                shapes.append(self.guess_dynamic_shape_object(*[o[i] for o in objs]))
+                shapes.append(
+                    self.guess_dynamic_shape_object(*[o[i] for o in objs], auto=auto, msg=msg)
+                )
             return tuple(shapes)
 
         if isinstance(obj, list):
@@ -298,7 +657,9 @@ class ModelInputs:
             ), f"Unexpected variety of list lengths {kl}{msg() if msg else ''}"
             shapes = []
             for i in range(kl.pop()):
-                shapes.append(self.guess_dynamic_shape_object(*[o[i] for o in objs]))
+                shapes.append(
+                    self.guess_dynamic_shape_object(*[o[i] for o in objs], auto=auto, msg=msg)
+                )
             return shapes
 
         if isinstance(obj, dict):
@@ -308,9 +669,33 @@ class ModelInputs:
             ), f"Unexpected variety of dict lengths {kl}{msg() if msg else ''}"
             shapes = {}
             for i in obj:
-                shapes[i] = self.guess_dynamic_shape_object(*[o[i] for o in objs])
+                shapes[i] = self.guess_dynamic_shape_object(
+                    *[o[i] for o in objs], auto=auto, msg=msg
+                )
             return shapes
 
+        if obj.__class__ in torch.utils._pytree.SUPPORTED_NODES:
+            kcl = set(o.__class__ for o in objs)
+            assert len(kcl) == 1, (
+                f"All instances of argument {i} are not of the same class but {kcl}, "
+                f"types should be the same."
+            )
+            col_args = [flatten_unflatten_for_dynamic_shapes(o) for o in objs]
+            kc = set(len(o) for o in col_args)
+            assert len(kc) == 1, (
+                f"All instances of type {kcl.pop()} are not serialized into the same number "
+                f"of arguments, it should be the same."
+            )
+            values = []
+            for i in range(kc.pop()):
+                values.append(
+                    self.guess_dynamic_shape_object(
+                        *[ca[i] for ca in col_args], auto=auto, msg=msg
+                    )
+                )
+            return values
+
+        # In case DynamicCache is not registered.
         if obj.__class__.__name__ == "DynamicCache":
             kc = set(len(o.key_cache) for o in objs)
             assert (
@@ -323,34 +708,39 @@ class ModelInputs:
             key_cache = []
             for i in range(kc.pop()):
                 key_cache.append(
-                    self.guess_dynamic_dimensions(*[o.key_cache[i] for o in objs])
+                    self.guess_dynamic_dimensions(*[o.key_cache[i] for o in objs], auto=auto)
                 )
             value_cache = []
             for i in range(vc.pop()):
                 value_cache.append(
-                    self.guess_dynamic_dimensions(*[o.value_cache[i] for o in objs])
+                    self.guess_dynamic_dimensions(*[o.value_cache[i] for o in objs], auto=auto)
                 )
             return [key_cache, value_cache]
 
         raise NotImplementedError(
             f"Unable to build dynamic shapes for type {set_types.pop()}: "
-            f"{string_type(objs)}{msg() if msg else ''} in {self.module_name_type}"
+            f"{string_type(objs)}{msg() if msg else ''} in {self.module_name_type}, "
+            f"this object needs serialization function to be registered."
         )
 
-    def guess_dynamic_shapes(
-        self,
-    ) -> Tuple[Tuple[Any, ...], Dict[str, Any]]:
+    def guess_dynamic_shapes(self, auto: bool = False) -> DYNAMIC_SHAPES:
         """
         Guesses the dynamic shapes for that module from two execution.
         If there is only one execution, then that would be static dimensions.
+
+        :param auto: if auto is True, use ``torch.export.Dim.AUTO`` for any
+            dimension if the number of inputs is one
         """
         if len(self.inputs) == 0:
             # No inputs, unable to guess.
             return (tuple(), {})
         if len(self.inputs) == 1:
             # No dynamic shapes.
-            return tuple(self.guess_dynamic_shape_object(a) for a in self.inputs[0][0]), {
-                k: self.guess_dynamic_shape_object(v) for k, v in self.inputs[0][1].items()
+            return tuple(
+                self.guess_dynamic_shape_object(a, auto=auto) for a in self.inputs[0][0]
+            ), {
+                k: self.guess_dynamic_shape_object(v, auto=auto)
+                for k, v in self.inputs[0][1].items()
             }
 
         # Otherwise.
@@ -365,7 +755,9 @@ class ModelInputs:
         for i in range(s1.pop()):
             objs = [_[0][i] for _ in self.inputs]
             args.append(
-                self.guess_dynamic_shape_object(*objs, msg=lambda i=i: f" failing input {i}")
+                self.guess_dynamic_shape_object(
+                    *objs, auto=auto, msg=lambda i=i: f" failing input {i}"
+                )
             )
         names = s2.pop()
         for name in names:
@@ -377,7 +769,7 @@ class ModelInputs:
 
             objs = [_[1][name] for _ in self.inputs]
             kwargs[name] = self.guess_dynamic_shape_object(
-                *objs, msg=lambda name=name: f" failing input {name!r}"
+                *objs, auto=auto, msg=lambda name=name: f" failing input {name!r}"
             )
         return tuple(args), kwargs
 
@@ -386,7 +778,7 @@ class ModelInputs:
         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 +826,22 @@ class ModelInputs:
                 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[Union[int, 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 []
+            dyn_shapes = self.guess_dynamic_shapes()
+        return [
+            CoupleInputsDynamicShapes(*i, dyn_shapes).invalid_dimensions_for_export()
+            for i in self.inputs
+        ]