onnx-diagnostic 0.8.2__py3-none-any.whl → 0.8.4__py3-none-any.whl

onnx_diagnostic/helpers/ort_session.py

@@ -1,7 +1,6 @@
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 import onnx
 import numpy as np
-import numpy.typing as npt
 import torch
 from torch._C import _from_dlpack
 import onnxruntime
@@ -16,6 +15,7 @@ from .torch_helper import torch_dtype_to_onnx_dtype
 
 
 DEVICES = {-1: ORTC.OrtDevice(ORTC.OrtDevice.cpu(), ORTC.OrtDevice.default_memory(), 0)}
+TensorLike = Union[np.ndarray, torch.Tensor]
 
 
 class _InferenceSession:
@@ -108,7 +108,10 @@ class _InferenceSession:
                     session_options,
                     providers=providers,
                 )
-            except onnxruntime.capi.onnxruntime_pybind11_state.Fail as e:
+            except (
+                onnxruntime.capi.onnxruntime_pybind11_state.Fail,
+                onnxruntime.capi.onnxruntime_pybind11_state.InvalidGraph,
+            ) as e:
                 if isinstance(sess, onnx.ModelProto):
                     debug_path = "_debug_InferenceSession_last_failure.onnx"
                     onnx.save(
@@ -134,7 +137,13 @@ class _InferenceSession:
 
         self.sess = sess
         self.input_names = [i.name for i in sess.get_inputs()]
+        assert (
+            "" not in self.input_names
+        ), f"Input name cannot be empty but input_names={self.input_names}"
         self.output_names = [i.name for i in sess.get_outputs()]
+        assert (
+            "" not in self.input_names
+        ), f"Output name cannot be empty but output_names={self.output_names}"
         self.input_shapes = [i.shape for i in sess.get_inputs()]
         self.output_shapes = [i.shape for i in sess.get_outputs()]
         self.input_types = [i.type for i in sess.get_inputs()]
@@ -234,16 +243,16 @@ class InferenceSessionForNumpy(_InferenceSession):
         )
 
     def run(
-        self, output_names: Optional[List[str]], feeds: Dict[str, npt.ArrayLike]
-    ) -> List[Optional[npt.ArrayLike]]:
+        self, output_names: Optional[List[str]], feeds: Dict[str, TensorLike]
+    ) -> List[Optional[TensorLike]]:
         """Calls :meth:`onnxruntime.InferenceSession.run`."""
         # sess.run does not support blfoat16
         # res = self.sess.run(output_names, feeds)
         return self._post_process_inplace(list(self.run_dlpack(output_names, feeds)))
 
     def run_dlpack(
-        self, output_names: Optional[List[str]], feeds: Dict[str, npt.ArrayLike]
-    ) -> Tuple[Optional[npt.ArrayLike], ...]:
+        self, output_names: Optional[List[str]], feeds: Dict[str, TensorLike]
+    ) -> Tuple[Optional[TensorLike], ...]:
         """
         Same as :meth:`onnxruntime.InferenceSession.run` except that
         feeds is a dictionary of :class:`np.ndarray`.
@@ -280,13 +289,13 @@ class InferenceSessionForNumpy(_InferenceSession):
     def _ortvalues_to_numpy_tensor(
         self,
         ortvalues: Union[List[ORTC.OrtValue], ORTC.OrtValueVector],
-    ) -> Tuple[Optional[npt.ArrayLike], ...]:
+    ) -> Tuple[Optional[TensorLike], ...]:
         if len(ortvalues) == 0:
             return tuple()
 
         if self.nvtx:
             self.torch.cuda.nvtx.range_push("_ortvalues_to_numpy_tensor")
-        res: List[Optional[npt.ArrayLike]] = []  # noqa: F823
+        res: List[Optional[TensorLike]] = []  # noqa: F823
         for i in range(len(ortvalues)):
             if not ortvalues[i].has_value():
                 res.append(None)
@@ -338,6 +347,7 @@ class InferenceSessionForTorch(_InferenceSession):
     :param optimized_model_filepath:  see :class:`onnxruntime.SessionOptions`
     :param disable_aot_function_inlining:  see :class:`onnxruntime.SessionOptions`
     :param use_training_api: use onnxruntime-traning API
+    :param cpu_output: if True, force the outputs to be on CPU
     """
 
     def __init__(
@@ -353,6 +363,7 @@ class InferenceSessionForTorch(_InferenceSession):
         optimized_model_filepath: Optional[str] = None,
         disable_aot_function_inlining: Optional[bool] = None,
         use_training_api: Optional[bool] = None,
+        cpu_outputs: bool = False,
     ):
         super().__init__(
             sess,
@@ -367,6 +378,7 @@ class InferenceSessionForTorch(_InferenceSession):
             disable_aot_function_inlining=disable_aot_function_inlining,
             use_training_api=use_training_api,
         )
+        self.cpu_outputs = cpu_outputs
 
     def _get_ortvalues_from_torch_tensors(
         self, tensors: Tuple[torch.Tensor, ...], n_outputs: int
@@ -490,23 +502,37 @@ class InferenceSessionForTorch(_InferenceSession):
         feeds is a dictionary of :class:`torch.Tensor`.
         The output device is CPU even if the outputs are on CUDA.
         """
-        new_feeds = {}
+        input_names = []
+        values = ORTC.OrtValueVector()
+        device = -1
         for k, v in feeds.items():
+            assert k != "", f"Input cannot be empty but feeds names={list(feeds)}"
+            device = max(device, v.get_device())
             assert hasattr(v, "__dlpack__"), f"class {type(v)} should be serialized"
             if not v.is_contiguous():
                 v = v.contiguous()
             if v.dtype == torch.bool:
-                # It does not work with dlpack
-                # unless onnxruntime updates the version it is using.
-                new_feeds[k] = ORTC.OrtValue.ortvalue_from_numpy_with_onnx_type(
-                    v.detach().numpy(), onnx.TensorProto.BOOL
-                )
+                v = v.to(torch.uint8)
+                v = ORTC.OrtValue.from_dlpack(v.__dlpack__(), True)
             else:
-                new_feeds[k] = ORTC.OrtValue.from_dlpack(v.__dlpack__(), False)
+                v = ORTC.OrtValue.from_dlpack(v.detach().__dlpack__(), False)
+            input_names.append(k)
+            values.push_back(v)
         if self.nvtx:
-            self.torch.cuda.nvtx.range_push("run_with_ort_values")
-        ort_outputs = self.sess._sess.run_with_ort_values(
-            new_feeds, output_names or self.output_names, self.run_options
+            self.torch.cuda.nvtx.range_push("run_with_ortvaluevector")
+
+        # ort_outputs = self.sess._sess.run_with_ort_values(
+        #    new_feeds, output_names or self.output_names, self.run_options
+        # )
+        ort_outputs = ORTC.OrtValueVector()
+        out_names = output_names or self.output_names
+        self.sess._sess.run_with_ortvaluevector(
+            self.run_options,
+            input_names,
+            values,
+            out_names,
+            ort_outputs,
+            [DEVICES[-1 if self.cpu_outputs else device] for o in out_names],
         )
         if self.nvtx:
             self.torch.cuda.nvtx.range_pop()
@@ -530,7 +556,7 @@ def investigate_onnxruntime_issue(
         Union[str, Callable[[onnx.ModelProto], onnxruntime.InferenceSession]]
     ] = None,
     # if model needs to be run.
-    feeds: Optional[Union[Dict[str, torch.Tensor], Dict[str, npt.ArrayLike]]] = None,
+    feeds: Optional[Union[Dict[str, torch.Tensor], Dict[str, TensorLike]]] = None,
     verbose: int = 0,
     dump_filename: Optional[str] = None,
     infer_shapes: bool = True,

onnx_diagnostic/helpers/torch_fx_graph_helper.py

@@ -0,0 +1,164 @@
+from typing import Any, Dict, Optional, Tuple
+import torch
+from .helper import string_type
+
+
+def validate_fx_tensor(
+    node: torch.fx.Node, tensor: torch.Tensor, expected_shape: Tuple[Any, ...]
+) -> None:
+    """
+    Validates the shape of tensor is expected.
+
+    :param node: node
+    :param tensor: tensor
+    :param expected_shape: expected shape
+    """
+    assert len(tensor.shape) == len(expected_shape), (
+        f"Shape mismatch, got {tensor.shape} expected {expected_shape}, "
+        f"node.name={node.name!r}, node.target={getattr(node, 'target', None)}, "
+        f"node.args={node.args}, node.kwargs={node.kwargs}, "
+        f"node.meta={node.meta}"
+    )
+    for a, b in zip(tensor.shape, expected_shape):
+        assert not isinstance(b, int) or a == b or {a, b} == {0, 1}, (
+            f"Dimension mismatch, got {tensor.shape} expected {expected_shape}, "
+            f"node.name={node.name!r}, node.target={getattr(node, 'target', None)}, "
+            f"node.args={node.args}, node.kwargs={node.kwargs}, "
+            f"node.meta={node.meta}"
+        )
+
+
+def validate_fx_outputs(node: torch.fx.Node, outputs: Tuple[Any, ...]) -> None:
+    """
+    Validates the outputs of a node using metadata stored in the node.
+
+    :param node: node
+    :param outputs: outputs
+    """
+    if "val" not in node.meta:
+        return
+    if isinstance(outputs, torch.Tensor):
+        validate_fx_tensor(node, outputs, node.meta["val"].shape)
+        return
+    if isinstance(outputs, (tuple, list)):
+        assert isinstance(node.meta["val"], (list, tuple)), (
+            f"Unexpected type {string_type(node.meta['val'])} for node.meta['val'], "
+            f"node.name={node.name!r}, node.target={getattr(node, 'target', None)}, "
+            f"node.args={node.args}, node.kwargs={node.kwargs}, "
+            f"node.meta={node.meta}"
+        )
+        assert len(outputs) == len(node.meta["val"]), (
+            f"Length mismatch, got {len(outputs)} expected {len(node.meta['val'])}, "
+            f"node.name={node.name!r}, node.target={getattr(node, 'target', None)}, "
+            f"node.args={node.args}, node.kwargs={node.kwargs}, "
+            f"node.meta={node.meta}"
+        )
+        for a, b in zip(outputs, node.meta["val"]):
+            validate_fx_tensor(node, a, b.shape)
+        return
+    if isinstance(outputs, int):
+        assert (
+            isinstance(node.meta["val"], (torch.SymInt, torch.SymBool, torch.SymFloat))
+            or outputs == node.meta["val"]
+        ), (
+            f"Int mismatch, got {outputs} expected {node.meta['val']}, "
+            f"node.name={node.name!r}, node.target={getattr(node, 'target', None)}, "
+            f"node.args={node.args}, node.kwargs={node.kwargs}, "
+            f"node.meta={node.meta}"
+        )
+        return
+    if outputs is None:
+        assert node.meta["val"] is None, (
+            f"None mismatch, got {outputs} expected {node.meta['val']}, "
+            f"node.name={node.name!r}, node.target={getattr(node, 'target', None)}, "
+            f"node.args={node.args}, node.kwargs={node.kwargs}, "
+            f"node.meta={node.meta}"
+        )
+        return
+    raise NotImplementedError(
+        f"Validation for output type {type(outputs)} is not implemented, "
+        f"node.name={node.name!r}, node.target={getattr(node, 'target', None)}, "
+        f"node.args={node.args}, node.kwargs={node.kwargs}, "
+        f"node.meta={node.meta}"
+    )
+
+
+def run_fx_node(
+    node: torch.fx.Node, args: Tuple[Any, ...], kwargs: Optional[Dict[str, Any]] = None
+) -> Tuple[Any, ...]:
+    """
+    Executes a node
+
+    :param node: runs a node
+    :param args: unnamed inputs to the node
+    :param kwargs: named inputs to the node
+    :return: results
+    """
+    if node.op == "output":
+        assert len(args) == 1 and not kwargs, (
+            f"Unexpected inputs: args={string_type(args, limit=20)} "
+            f"kwargs={string_type(kwargs, limit=20)}"
+        )
+        return args
+    if node.op == "call_function":
+        assert callable(node.target), f"{node.target!r} not callable in node {node!r}"
+        for a, ea in zip(args, node.args):
+            if isinstance(a, torch.Tensor) and hasattr(ea, "meta") and "val" in ea.meta:
+                ta = ea.meta["val"]
+                assert (
+                    isinstance(ta, torch.Tensor)
+                    and len(a.shape) == len(ta.shape)
+                    and a.dtype == ta.dtype
+                ), (
+                    f"Unable to run node {node!r}, target={node.target!r}, "
+                    f"node.args={node.args!r}, node.kwargs={node.kwargs!r}, "
+                    f"args={string_type(args, with_shape=True, with_device=True)}, "
+                    f"kwargs={string_type(kwargs, with_shape=True, with_device=True)}"
+                )
+        try:
+            outputs = node.target(*args, **(kwargs or {}))
+        except RuntimeError as e:
+            raise RuntimeError(
+                f"Unable to run node {node!r}, target={node.target!r}, "
+                f"args={string_type(args, with_shape=True, with_device=True)}, "
+                f"kwargs={string_type(kwargs, with_shape=True, with_device=True)}"
+            ) from e
+        validate_fx_outputs(node, outputs)
+        return outputs
+    raise NotImplementedError(
+        f"node.op={node.op!r} is not implemented, node.name={node.name!r}"
+    )
+
+
+def _pick_result(torch_results: Dict[str, Any], ref: Any) -> Any:
+    "See :func:`prepare_args_kwargs`."
+    if isinstance(ref, torch.fx.Node):
+        return torch_results[ref.name]
+    if isinstance(ref, list):
+        return [_pick_result(torch_results, n) for n in ref]
+    if isinstance(ref, tuple):
+        return tuple(_pick_result(torch_results, n) for n in ref)
+    if isinstance(ref, dict):
+        return {k: _pick_result(torch_results, v) for k, v in ref.items()}
+    if isinstance(ref, (bool, int, float, str, torch.device, torch.dtype)):
+        return ref
+    if ref is None:
+        return None
+    if isinstance(ref, torch.layout):
+        return ref
+    raise NotImplementedError(f"Unable to process args type {type(ref)}")
+
+
+def prepare_args_kwargs(
+    torch_results: Dict[str, Any], node: torch.fx.Node
+) -> Tuple[Tuple[Any, ...], Dict[str, Any]]:
+    """
+    Prepares args and kwargs before executing a fx node.
+
+    :param torch_results: existing results
+    :param node: node to execute
+    :return: new args and kwargs
+    """
+    new_args = _pick_result(torch_results, node.args)
+    new_kwargs = _pick_result(torch_results, node.kwargs)
+    return new_args, new_kwargs

onnx_diagnostic/helpers/torch_helper.py

@@ -1,6 +1,7 @@
 import contextlib
 import ctypes
 import inspect
+import math
 import os
 import sys
 import warnings
@@ -30,9 +31,7 @@ from .onnx_helper import (
 
 
 def proto_from_tensor(
-    arr: "torch.Tensor",  # noqa: F821
-    name: Optional[str] = None,
-    verbose: int = 0,
+    arr: torch.Tensor, name: Optional[str] = None, verbose: int = 0
 ) -> onnx.TensorProto:
     """
     Converts a torch Tensor into a TensorProto.
@@ -98,7 +97,7 @@ def proto_from_tensor(
     return tensor
 
 
-def onnx_dtype_to_torch_dtype(itype: int) -> "torch.dtype":  # noqa: F821
+def onnx_dtype_to_torch_dtype(itype: int) -> torch.dtype:
     """
     Converts an onnx type into a torch dtype.
 
@@ -140,7 +139,16 @@ def onnx_dtype_to_torch_dtype(itype: int) -> "torch.dtype":  # noqa: F821
     )
 
 
-def torch_dtype_to_onnx_dtype(to: "torch.dtype") -> int:  # noqa: F821
+_TYPENAME = dict(
+    FLOAT=onnx.TensorProto.FLOAT,
+    INT64=onnx.TensorProto.INT64,
+    INT32=onnx.TensorProto.INT32,
+    FLOAT16=onnx.TensorProto.FLOAT16,
+    BFLOAT16=onnx.TensorProto.BFLOAT16,
+)
+
+
+def torch_dtype_to_onnx_dtype(to: torch.dtype) -> int:
     """
     Converts a torch dtype into a onnx element type.
 
@@ -183,7 +191,13 @@ def torch_dtype_to_onnx_dtype(to: "torch.dtype") -> int:  # noqa: F821
         return onnx.TensorProto.COMPLEX64
     if to == torch.complex128:
         return onnx.TensorProto.COMPLEX128
-    raise NotImplementedError(f"Unable to convert torch dtype {to!r} to onnx dtype.")
+    # SymbolicTensor
+    sto = str(to)
+    if sto in _TYPENAME:
+        return _TYPENAME[sto]
+    raise NotImplementedError(
+        f"Unable to convert torch dtype {to!r} ({type(to)}) to onnx dtype."
+    )
 
 
 def _forward_(
@@ -483,7 +497,7 @@ def is_torchdynamo_exporting() -> bool:
             return False
 
 
-def to_numpy(tensor: "torch.Tensor") -> np.ndarray:  # noqa: F821
+def to_numpy(tensor: torch.Tensor) -> np.ndarray:
     """Converts a :class:`torch.Tensor` to :class:`numpy.ndarray`."""
     try:
         return tensor.detach().cpu().numpy()
@@ -498,6 +512,21 @@ def to_numpy(tensor: "torch.Tensor") -> np.ndarray:  # noqa: F821
     return tensor.detach().to(torch.float32).cpu().numpy().astype(conv[tensor.dtype])
 
 
+def from_numpy(tensor: np.ndarray) -> torch.Tensor:
+    """Converts a :class:`numpy.ndarray` to :class:`torch.Tensor`."""
+    try:
+        return torch.from_numpy(tensor)
+    except TypeError:
+        # We try with ml_dtypes
+        pass
+
+    import ml_dtypes
+
+    conv = {ml_dtypes.bfloat16: torch.bfloat16}
+    assert tensor.dtype in conv, f"Unsupported type {tensor.dtype}, not in {conv}"
+    return torch.from_numpy(tensor.astype(torch.float32)).to(conv[tensor.dtype])
+
+
 def replace_string_by_dynamic(dynamic_shapes: Any) -> Any:
     """Replaces strings by ``torch.export.Dim.DYNAMIC``."""
     import torch
@@ -797,7 +826,8 @@ def torch_deepcopy(value: Any) -> Any:
     if isinstance(value, tuple):
         return tuple(torch_deepcopy(v) for v in value)
     if isinstance(value, list):
-        return [torch_deepcopy(v) for v in value]
+        if type(value) is list:
+            return [torch_deepcopy(v) for v in value]
     if isinstance(value, set):
         return {torch_deepcopy(v) for v in value}
     if isinstance(value, dict):
@@ -990,3 +1020,96 @@ def get_weight_type(model: torch.nn.Module) -> torch.dtype:
             counts[dt] += 1
     final = max(list(counts.items()))
     return final[0]
+
+
+def closest_factor_pair(n: int):
+    """Tries to find ``a, b`` such as ``n == a * b``."""
+    assert n > 0, f"n={n} must be a positive integer"
+    start = math.isqrt(n)
+    for a in range(start, 0, -1):
+        if n % a == 0:
+            b = n // a
+            return a, b
+    return 1, n
+
+
+def study_discrepancies(
+    t1: torch.Tensor,
+    t2: torch.Tensor,
+    bins: int = 50,
+    figsize: Optional[Tuple[int, int]] = (15, 15),
+    title: Optional[str] = None,
+    name: Optional[str] = None,
+) -> "matplotlib.axes.Axes":  # noqa: F821
+    """
+    Computes different metrics for the discrepancies.
+    Returns graphs.
+
+    .. plot::
+        :include-source:
+
+        import torch
+        from onnx_diagnostic.helpers.torch_helper import study_discrepancies
+
+        t1 = torch.randn((512, 1024)) * 10
+        t2 = t1 + torch.randn((512, 1024))
+        study_discrepancies(t1, t2, title="Random noise")
+    """
+    assert t1.dtype == t2.dtype, f"Type mismatch {t1.dtype} != {t2.dtype}"
+    assert t1.shape == t2.shape, f"Shape mismatch {t1.shape} != {t2.shape}"
+    d1, d2 = (
+        (t1, t2) if t1.dtype == torch.float64 else (t1.to(torch.float32), t2.to(torch.float32))
+    )
+
+    d1 = d1.squeeze()
+    d2 = d2.squeeze()
+    if len(d1.shape) == 1:
+        new_shape = closest_factor_pair(d1.shape[0])
+        d1, d2 = d1.reshape(new_shape), d2.reshape(new_shape)
+    elif len(d1.shape) > 2:
+        new_shape = (-1, max(d1.shape))
+        d1, d2 = d1.reshape(new_shape), d2.reshape(new_shape)
+
+    import matplotlib.pyplot as plt
+
+    fig, ax = plt.subplots(3, 2, figsize=figsize)
+    vmin, vmax = d1.min().item(), d1.max().item()
+    ax[0, 0].imshow(d1.detach().cpu().numpy(), cmap="Greys", vmin=vmin, vmax=vmax)
+    ax[0, 0].set_title(
+        f"Color plot of the first tensor in\n[{vmin}, {vmax}]\n{t1.shape} -> {d1.shape}"
+    )
+
+    diff = d2 - d1
+    vmin, vmax = diff.min().item(), diff.max().item()
+    ax[0, 1].imshow(diff.detach().cpu().numpy(), cmap="seismic", vmin=vmin, vmax=vmax)
+    ax[0, 1].set_title(f"Color plot of the differences in \n[{vmin}, {vmax}]")
+
+    ax[1, 0].hist(d1.detach().cpu().numpy().ravel(), bins=bins)
+    ax[1, 0].set_title("Distribution of the first tensor")
+
+    ax[1, 1].hist(diff.detach().cpu().numpy().ravel(), bins=bins)
+    ax[1, 1].set_title("Distribution of the differences")
+
+    tf1 = d1.ravel()
+    td1 = diff.ravel()
+    ax[2, 1].plot(tf1.detach().cpu().numpy(), td1.detach().cpu().numpy(), ".")
+    ax[2, 1].set_title("Graph XY")
+    ax[2, 1].set_xlabel("First tensor values")
+    ax[2, 1].set_ylabel("Difference values")
+
+    if title:
+        fig.suptitle(title)
+    fig.tight_layout()
+    if name:
+        fig.savefig(name)
+    return ax
+
+
+def int_device_to_torch_device(device_id: int) -> torch.device:
+    """
+    Converts a device defined as an integer (coming from :meth:`torch.Tensor.get_device`)
+    into a ``torch.device``.
+    """
+    if device_id < 0:
+        return torch.device("cpu")
+    return torch.device("cuda", device_id)