onnx-diagnostic 0.8.9__py3-none-any.whl → 0.8.11__py3-none-any.whl

onnx_diagnostic/investigate/input_observer.py

@@ -0,0 +1,329 @@
+import contextlib
+import inspect
+from typing import Any, Callable, Sequence
+import torch
+
+
+def flatten_unflatten_for_dynamic_shapes(
+    obj: Any,
+    use_dict: bool = True,
+    change_function: Callable[[torch.Tensor], Any] | None = None,
+) -> Any:
+    """
+    Returns the object in a different structure similar to what
+    the definition of the dynamic shapes should use.
+
+    Args:
+        obj:
+            object from a custom class
+        use_dict:
+            closer to the original result but
+            :func:`torch.export.export` only considers the values,
+            the context gives the dictionary keys but it is not expressed
+            in the dynamic shapes, these specifications seems to be different
+            for the strict and non strict mode. It also preserves tuple.
+        change_function:
+            to modify the tensor in the structure itself,
+            like replace them by a shape
+
+    Returns:
+        the serialized object
+    """
+    if isinstance(obj, torch.Tensor):
+        return change_function(obj) if change_function else obj
+    flat, spec = torch.utils._pytree.tree_flatten(obj)
+    start = 0
+    end = 0
+    subtrees = []
+    for subspec in (spec.children() if hasattr(spec, "children") else spec.children_specs):
+        end += subspec.num_leaves
+        value = subspec.unflatten(flat[start:end])
+        value = flatten_unflatten_for_dynamic_shapes(
+            value, use_dict=use_dict, change_function=change_function
+        )
+        subtrees.append(value)
+        start = end
+    if use_dict:
+        if spec.type is dict:
+            # This is a dictionary.
+            return dict(zip(spec.context, subtrees))
+        if spec.type is tuple:
+            return tuple(subtrees)
+        if spec.type is list:
+            return list(subtrees)
+        if spec.type is None and not subtrees:
+            return None
+        if spec.context:
+            # This is a custom class with attributes.
+            # It is returned as a list.
+            return list(subtrees)
+        raise ValueError(
+            f"Unable to interpret spec type {spec.type} "
+            f"(type is {type(spec.type)}, context is {spec.context}), "
+            f"spec={spec}, subtrees={subtrees}"
+        )
+    # This is a list.
+    return subtrees
+
+
+def infer_dynamic_dimensions(shape_list: Sequence[tuple[int, ...]]) -> list[int]:
+    """
+    Returns the list of dynamic dimensions given a list of shapes
+    corresponding to the same tensor.
+
+    Args:
+        shape_list:
+            list of shapes, they must all have the same length
+
+    Returns:
+        list of dynamic dimensions
+    """
+    unique_ranks = {len(shape) for shape in shape_list}
+    torch._check(
+        len(unique_ranks) == 1, lambda: "all shapes in shape_list must have the same rank"
+    )
+    rank = unique_ranks.pop()
+    dynamic = []
+    for i in range(rank):
+        dims = [shape[i] for shape in shape_list]
+        if len(set(dims)) > 1:
+            dynamic.append(i)
+    return dynamic
+
+
+class InputObserverInfo:
+    def __init__(self, signature: inspect.Signature):
+        # pyrefly: ignore
+        self.inputs_specs: list[torch.utils._pytree.PyTreeSpec] = []
+        self.flat_inputs: list[list[torch.Tensor | None]] = []
+        # pyrefly: ignore
+        self.outputs_specs: list[torch.utils._pytree.PyTreeSpec] = []
+        self.flat_outputs: list[torch.Tensor | list[torch.Tensor]] = []
+        self.signature = signature
+
+        self._max_args: tuple[Any, torch.Tensor] | None = None
+        self._max_kwargs: dict[str, torch.Tensor] | None = None
+
+    def __len__(self) -> int:
+        return len(self.flat_inputs)
+
+    def add_inputs(self, args: tuple[Any, ...], kwargs: dict[str, Any]):
+        kwargs = {
+            k: v
+            for k, v in kwargs.items()
+            if v is not None and not isinstance(v, (int, float, bool))
+        }
+        flat_args, spec = torch.utils._pytree.tree_flatten((args, kwargs))
+        self.inputs_specs.append(spec)
+        cloned = [
+            (None if not isinstance(t, torch.Tensor) else t.clone().detach())
+            for t in flat_args
+        ]
+        self.flat_inputs.append(cloned)
+
+        cloned_args, cloned_kwargs = torch.utils._pytree.tree_unflatten(cloned, spec)
+        if self._max_args is None or len(cloned_args) > len(self._max_args):
+            self._max_args = cloned_args
+        if self._max_kwargs is None or len(cloned_kwargs) > len(self._max_kwargs):
+            self._max_kwargs = cloned_kwargs
+
+    def add_outputs(self, res: torch.Tensor | tuple[torch.Tensor, ...]):
+        flat_res, spec = torch.utils._pytree.tree_flatten(res)
+        self.outputs_specs.append(spec)
+        self.flat_outputs.append([t.clone().detach() for t in flat_res])
+
+    def build_inputs_completed_with_none_values(self) -> list[list[torch.Tensor]]:
+        # Let's compute the sizes of each independently.
+        if not self.flat_inputs or self._max_args is None or self._max_kwargs is None:
+            raise RuntimeError("No inputs were captured.")
+        arg_sizes = [len(torch.utils._pytree.tree_flatten(a)[0]) for a in self._max_args]
+        kwarg_sizes = {
+            k: len(torch.utils._pytree.tree_flatten(v)[0]) for k, v in self._max_kwargs.items()
+        }
+
+        # Let's reprocess everything.
+        captured_inputs: dict[int | str, int] = {}
+        new_flat_inputs = []
+        for args_kwargs, spec in zip(self.flat_inputs, self.inputs_specs):
+            args, kwargs = torch.utils._pytree.tree_unflatten(args_kwargs, spec)
+            if len(set(kwargs) | set(self._max_kwargs)) > len(self._max_kwargs):
+                raise RuntimeError(
+                    "At least one call to the observed model "
+                    "must contain all the named arguments."
+                )
+            flat = []
+            for i in range(len(self._max_args)):
+                if i < len(args):
+                    ts = torch.utils._pytree.tree_flatten(args[i])[0]
+                    if i in captured_inputs and captured_inputs[i] != len(ts):
+                        raise RuntimeError(
+                            f"Positional argument {i} has {len(ts)} tensors "
+                            f"but previously got {captured_inputs[i]} tensors. "
+                            f"Inference is impossible in that case."
+                        )
+                    captured_inputs[i] = len(ts)
+                    flat.extend(ts)
+                else:
+                    flat.extend([None for _ in range(arg_sizes[i])])
+            for k in self._max_kwargs:
+                if k in kwargs:
+                    ts = torch.utils._pytree.tree_flatten(kwargs[k])[0]
+                    if k in captured_inputs and captured_inputs[k] != len(ts):
+                        raise RuntimeError(
+                            f"Named argument {k!r} has {len(ts)} tensors "
+                            f"but previously got {captured_inputs[k]} tensors. "
+                            f"Inference is impossible in that case."
+                        )
+                    captured_inputs[k] = len(ts)
+                    flat.extend(ts)
+                else:
+                    flat.extend([None for _ in range(kwarg_sizes[k])])
+            new_flat_inputs.append(flat)
+        return new_flat_inputs
+
+    def infer_dynamic_shapes(self) -> tuple[dict[int, Any], ...] | dict[str, dict[int, Any]]:
+        flat_inputs = self.build_inputs_completed_with_none_values()
+        # This is already checked by build_inputs_completed_with_none_values
+        # but this is not always well captured by tools checking types.
+        assert self._max_args is not None and self._max_kwargs is not None
+        if len({len(flat) for flat in flat_inputs}) != 1:
+            raise NotImplementedError(
+                "infer_dynamic_shapes is not implemented "
+                "when the number of input tensors are not the same."
+            )
+        shape_lists = [
+            [(None if t is None else t.shape) for t in tensors] for tensors in flat_inputs
+        ]
+        n_tensors = len(shape_lists[0])
+        dynamic_shapes = [
+            infer_dynamic_dimensions(
+                [s for s in [shapes[index] for shapes in shape_lists] if s is not None]
+            )
+            for index in range(n_tensors)
+        ]
+        cst = torch.export.Dim.DYNAMIC
+        flat_dynamic_shapes = [dict.fromkeys(dims, cst) for dims in dynamic_shapes]
+        if len(flat_dynamic_shapes) == len(self._max_args) + len(self._max_kwargs):
+            # It means forward method is called with tensors only.
+            if not self._max_kwargs:
+                # only positional arguments
+                return tuple(flat_dynamic_shapes)
+            if not self._max_args:
+                # only named arguments
+                return dict(zip(list(self._max_kwargs), flat_dynamic_shapes))
+            # positional arguments needs to be moved to the named arguments
+            n_args = len(self._max_args)
+            pos_names = list(self.signature.parameters)[:n_args]
+            return {
+                **dict(zip(pos_names, flat_dynamic_shapes[:n_args])),
+                **dict(zip(list(self._max_kwargs), flat_dynamic_shapes[n_args:])),
+            }
+
+        # nested types, here comes the fun part because the shapes cannot be unflattened,
+        # custom classes must appear in their flattened shape.
+        # This does not work in all cases but every time every available argument is flattened
+        # with the same number of tensors. The function does not check
+        # if that assumption is true.
+        flat_inputs, _max_spec = torch.utils._pytree.tree_flatten(
+            (self._max_args, self._max_kwargs)
+        )
+        torch._check(
+            len(flat_inputs) == len(flat_dynamic_shapes),
+            (
+                f"Length mismatch len(flat_inputs)={len(flat_inputs)}, "
+                f"len(flat_dynamic_shapes)={len(flat_dynamic_shapes)}"
+            ),
+        )
+        mapping = {id(t): shape for t, shape in zip(flat_inputs, flat_dynamic_shapes)}
+        ds_args, ds_kwargs = flatten_unflatten_for_dynamic_shapes(
+            (self._max_args, self._max_kwargs), change_function=lambda t: mapping[id(t)]
+        )
+        if not ds_kwargs:
+            return tuple(ds_args)
+        if not ds_args:
+            return tuple(ds_kwargs)
+        pos_names = list(self.signature.parameters)[: len(ds_args)]
+        return {**dict(zip(pos_names, ds_args)), **ds_kwargs}
+
+    def infer_arguments(
+        self, index: int | None = None
+    ) -> tuple[torch.Tensor, ...] | dict[str, torch.Tensor]:
+        # This is already checked by build_inputs_completed_with_none_values
+        # but this is not always well captured by tools checking types.
+        assert self._max_args is not None and self._max_kwargs is not None
+        candidate = None
+        if index is None:
+            for i, (args_kwargs, spec) in enumerate(zip(self.flat_inputs, self.inputs_specs)):
+                args, kwargs = torch.utils._pytree.tree_unflatten(args_kwargs, spec)
+                if len(args) == len(self._max_args) and len(kwargs) == len(self._max_kwargs):
+                    index = i
+                    candidate = args, kwargs
+                    break
+        if index is not None:
+            # found one available set.
+            args, kwargs = candidate or torch.utils._pytree.tree_unflatten(
+                self.flat_inputs[index], self.inputs_specs[index]
+            )
+            if not kwargs:
+                return args
+            if not args:
+                return kwargs
+            # We need to move args to kwargs
+            pos_names = list(self.signature.parameters)[: len(args)]
+            return {**dict(zip(pos_names, args)), **kwargs}
+
+        raise NotImplementedError(
+            "We could not find a good set of inputs/outputs. "
+            "We need to replace none by empty tensors."
+        )
+
+
+class InputObserver:
+    def __init__(self, store_n_calls: int = 3):
+        self.store_n_calls = store_n_calls
+        self.info: InputObserverInfo | None = None
+
+    def _forward_captured(self, *args, _captured_forward=None, **kwargs):
+        assert _captured_forward is not None, "_captured_forward cannot be None"
+        assert self.info is not None, "info cannot be None"
+        n_stored = len(self.info)
+        if n_stored < self.store_n_calls:
+            self.info.add_inputs(args, kwargs)
+        res = _captured_forward(*args, **kwargs)
+        if n_stored < self.store_n_calls:
+            self.info.add_outputs(res)
+        return res
+
+    @contextlib.contextmanager
+    def __call__(self, model: torch.nn.Module):
+        if self.info is not None:
+            raise RuntimeError(
+                "This class was already used to capture a model. Please create a new one."
+            )
+        self.info = InputObserverInfo(signature=inspect.signature(model.forward))
+        forward_method = model.forward
+        model.forward = (
+            lambda *args, _captured_forward=forward_method, **kwargs: self._forward_captured(
+                *args, _captured_forward=_captured_forward, **kwargs
+            )
+        )
+        try:
+            yield self
+        finally:
+            model.forward = forward_method
+
+    def _check_captured(self):
+        if self.info is None:
+            raise RuntimeError("No inputs were captured.")
+
+    def infer_dynamic_shapes(self) -> tuple[dict[int, Any], ...] | dict[str, dict[int, Any]]:
+        self._check_captured()
+        assert self.info is not None  # missed by type checking
+        return self.info.infer_dynamic_shapes()
+
+    def infer_arguments(
+        self, index: int | None = None
+    ) -> tuple[torch.Tensor, ...] | dict[str, torch.Tensor]:
+        self._check_captured()
+        assert self.info is not None  # missed by type checking
+        return self.info.infer_arguments(index=index)

onnx_diagnostic/reference/evaluator.py

@@ -42,7 +42,6 @@ from .ops.op_slice import Slice_1, Slice_10
 from .ops.op_transpose_cast import Transpose2DCastFP16, Transpose2DCastFP32
 from .ops.op_tri_matrix import TriMatrix
 
-
 logger = getLogger("onnx-diagnostic-eval")
 
 

onnx_diagnostic/reference/ort_evaluator.py

@@ -34,7 +34,6 @@ from ..helpers.torch_helper import to_tensor
 from .report_results_comparison import ReportResultComparison
 from .evaluator import ExtendedReferenceEvaluator
 
-
 PROTO = (FunctionProto, ModelProto, GraphProto, NodeProto)
 Proto = Union[FunctionProto, ModelProto, GraphProto, NodeProto]
 

onnx_diagnostic/reference/report_results_comparison.py

@@ -1,5 +1,4 @@
-from typing import Any, Dict, List, Tuple, Union
-
+from typing import Any, Dict, List, Set, Tuple, Union
 
 ReportKeyNameType = Union[str, Tuple[str, int, str]]
 ReportKeyValueType = Tuple[int, Tuple[int, ...]]
@@ -14,6 +13,7 @@ class ReportResultComparison:
     :param tensors: tensor
     """
 
+    # pyrefly: ignore[unknown-name]
     def __init__(self, tensors: Dict[ReportKeyNameType, "torch.Tensor"]):  # noqa: F821
         from ..helpers.onnx_helper import dtype_to_tensor_dtype
         from ..helpers import max_diff, string_type
@@ -25,7 +25,9 @@ class ReportResultComparison:
         self.max_diff = max_diff
         self.tensors = tensors
         self._build_mapping()
+        self.unique_run_names: Set[str] = set()
 
+    # pyrefly: ignore[unknown-name]
     def key(self, tensor: "torch.Tensor") -> ReportKeyValueType:  # noqa: F821
         "Returns a key for a tensor, (onnx dtype, shape)."
         return self.dtype_to_tensor_dtype(tensor.dtype), tuple(map(int, tensor.shape))
@@ -59,12 +61,15 @@ class ReportResultComparison:
         for k, v in self.value.items():
             (i_run, run_name), ref_name = k
             d = dict(run_index=i_run, run_name=run_name, ref_name=ref_name)
+            # pyrefly: ignore[no-matching-overload]
             d.update(v)
             rows.append(d)
         return rows
 
     def report(
-        self, outputs: Dict[str, "torch.Tensor"]  # noqa: F821
+        self,
+        # pyrefly: ignore[unknown-name]
+        outputs: Dict[str, "torch.Tensor"],  # noqa: F821
     ) -> List[Tuple[Tuple[int, str], ReportKeyNameType, Dict[str, Union[float, str]]]]:
         """
         For every tensor in outputs, compares it to every tensor held by
@@ -79,6 +84,7 @@ class ReportResultComparison:
             key = self.key(tensor)
             if key not in self.mapping:
                 continue
+            # pyrefly: ignore[unknown-name]
             cache: Dict["torch.device", "torch.Tensor"] = {}  # noqa: F821, UP037
             for held_key in self.mapping[key]:
                 t2 = self.tensors[held_key]

onnx_diagnostic/reference/torch_evaluator.py

@@ -63,7 +63,7 @@ class TorchOnnxEvaluator:
     * `functions`: local functions
 
     The class is not multithreaded. `runtime_info` gets updated
-    by the the class. The list of available kernels is returned by function
+    by the class. The list of available kernels is returned by function
     :func:`onnx_diagnostic.reference.torch_evaluator.get_kernels`.
     Example:
 
@@ -494,8 +494,10 @@ class TorchOnnxEvaluator:
             r = self.runtime_info[k]
             r.set_value(
                 torch_ops.OpRunTensor(
+                    # pyrefly: ignore[missing-attribute]
                     v.to(self.CUDA) if not r.is_shape and self.on_cuda else v,
                     is_constant=False,
+                    # pyrefly: ignore[missing-attribute]
                     may_cpu=len(v.shape) == 1 and v.numel() < 8 and v.dtype == torch.int64,
                 )
             )
@@ -524,6 +526,7 @@ class TorchOnnxEvaluator:
                         f"for kernel {type(kernel)}."
                     )
                     for name, t in zip(kernel.output, res):
+                        # pyrefly: ignore[bad-argument-type]
                         self.runtime_info[name].set_value(t)
                     if self.verbose:
                         for name in kernel.output:
@@ -644,6 +647,7 @@ class TorchOnnxEvaluator:
                         f"for kernel {type(kernel)}."
                     )
                     for name, t in zip(kernel.output, res):
+                        # pyrefly: ignore[bad-argument-type]
                         self.runtime_info[name].set_value(t)
                 else:
                     assert isinstance(

onnx_diagnostic/reference/torch_ops/_op_run.py

@@ -1,7 +1,7 @@
 from typing import Any, Dict, List, Optional, Union, Tuple
 import onnx
 import torch
-from ...api import TensorLike
+from ...typing import TensorLike
 from ...helpers import string_type
 from ...helpers.torch_helper import to_tensor
 
@@ -149,7 +149,7 @@ class OpRunSequence(OpRunValue):
     ) -> "OpRunSequence":
         "Inserts a value at a given position."
         assert isinstance(tensor, OpRunTensor), f"Unexpected type {type(tensor)} for tensor"
-        new_seq = OpRunSequence()
+        new_seq = OpRunSequence()  # type: ignore[abstract]
         seq = self.sequence.copy()
         new_seq.sequence = seq
         if position is None:
@@ -314,9 +314,7 @@ class OpRunKernel:
 
 
 class OpRunFunction(OpRunKernel):
-    """
-    Defines a kernel based on a local functions.
-    """
+    """Defines a kernel based on a local functions."""
 
     def __init__(
         self,

onnx_diagnostic/reference/torch_ops/sequence_ops.py

@@ -46,7 +46,7 @@ class SequenceEmpty_11(OpRunOpSequence):
         )
 
     def run(self) -> OpRunSequence:
-        return OpRunSequence(dtype=self.dtype)
+        return OpRunSequence(dtype=self.dtype)  # type: ignore[abstract]
 
 
 class SequenceInsert_11(OpRunOpSequence):

onnx_diagnostic/tasks/feature_extraction.py

@@ -3,7 +3,6 @@ import torch
 from ..helpers.config_helper import update_config, check_hasattr
 from ..helpers.cache_helper import make_dynamic_cache, make_encoder_decoder_cache
 
-
 __TASK__ = "feature-extraction"
 
 

onnx_diagnostic/torch_export_patches/__init__.py

@@ -4,7 +4,6 @@ from .onnx_export_errors import (
 )
 from .patch_module import torch_export_rewrite
 
-
 # bypass_export_some_errors is the first name given to the patches.
 bypass_export_some_errors = torch_export_patches  # type: ignore
 

onnx_diagnostic/torch_export_patches/onnx_export_serialization.py

@@ -305,7 +305,7 @@ def serialization_functions(
 
 
 def unregister_class_serialization(cls: type, verbose: int = 0):
-    """Undo the registration."""
+    """Undo the registration for a class."""
     # torch.utils._pytree._deregister_pytree_flatten_spec(cls)
     if cls in torch.fx._pytree.SUPPORTED_NODES:
         del torch.fx._pytree.SUPPORTED_NODES[cls]
@@ -333,6 +333,10 @@ def unregister_class_serialization(cls: type, verbose: int = 0):
 
 
 def unregister_cache_serialization(undo: Dict[str, bool], verbose: int = 0):
+    """
+    Undo the registration made by
+    :func:`onnx_diagnostic.torch_export_patches.onnx_export_serialization.register_cache_serialization`.
+    """
     cls_ensemble = {DynamicCache, EncoderDecoderCache} | set(undo)
     for cls in cls_ensemble:
         if undo.get(cls.__name__, False):

onnx_diagnostic/torch_export_patches/patch_module.py

@@ -986,7 +986,7 @@ def torch_export_rewrite(
             name = me.__qualname__
             spl = name.split(".")
             if len(spl) == 1:
-                # This a function
+                # This is a function
                 module = me.__module__
                 if module in me.__globals__:
                     mod = me.__globals__[module]

onnx_diagnostic/torch_export_patches/patches/_patch_transformers_rotary_embedding.py

@@ -7,10 +7,10 @@ import transformers
 
 def patched__compute_dynamic_ntk_parameters(
     config: Optional[transformers.PretrainedConfig] = None,
-    device: Optional["torch.device"] = None,
+    device: Optional[torch.device] = None,
     seq_len: Optional[int] = None,
     **rope_kwargs,
-) -> Tuple["torch.Tensor", float]:
+) -> Tuple[torch.Tensor, float]:
     """
     manual patch:
     ``[patch:transformers.modeling_rope_utils._compute_dynamic_ntk_parameters]``

onnx_diagnostic/torch_export_patches/patches/patch_torch.py

@@ -524,13 +524,16 @@ class patched_ShapeEnv:
 
             transmute_into_runtime_assert = False
 
+            backed_var_to_val = getattr(
+                self, "backed_var_to_val", getattr(self, "var_to_val", {})
+            )
             concrete_val = None
-            if not (expr.free_symbols <= self.var_to_val.keys()):
+            if not (expr.free_symbols <= backed_var_to_val.keys()):
                 # TODO: dedupe this with _maybe_evaluate_static
                 # Attempt to eliminate the unbacked SymInt
                 new_expr = self._maybe_evaluate_static(expr, unbacked_only=True)
                 assert new_expr is not None
-                if not (new_expr.free_symbols <= self.var_to_val.keys()):
+                if not (new_expr.free_symbols <= backed_var_to_val.keys()):
                     ok = False
 
                     # fallback_value is set when guard_or_true or guard_or_false are used.
@@ -541,17 +544,15 @@ class patched_ShapeEnv:
                     # oblivious_var_to_val will be defined iff we have sizes
                     # with DimDynamic.OBLIVIOUS_SIZE type.
                     # See https://github.com/pytorch/pytorch/issues/137100#issuecomment-2495778113
-                    var_to_val = getattr(
-                        self,
-                        "unbacked_var_to_val",
-                        getattr(self, "oblivious_var_to_val", False),
-                    )
                     if (
-                        var_to_val
-                        and not (correct_hint := orig_expr.xreplace(var_to_val)).free_symbols
+                        backed_var_to_val
+                        and getattr(self, "real_tensor_prop_unbacked_vals", True)
+                        and not (
+                            correct_hint := orig_expr.xreplace(backed_var_to_val)
+                        ).free_symbols
                         and not (
                             counterfactual_hint := orig_expr.xreplace(
-                                {k: max(2, v) for k, v in var_to_val.items()}
+                                {k: max(2, v) for k, v in backed_var_to_val.items()}
                             )
                         ).free_symbols
                         and correct_hint == counterfactual_hint
@@ -574,10 +575,10 @@ class patched_ShapeEnv:
                     # and if they pass we add a runtime assertions and continue.
                     if (
                         not ok
-                        and var_to_val
+                        and backed_var_to_val
                         and not (
-                            unsound_result := orig_expr.xreplace(var_to_val).xreplace(
-                                var_to_val
+                            unsound_result := orig_expr.xreplace(backed_var_to_val).xreplace(
+                                backed_var_to_val
                             )
                         ).free_symbols
                     ):