onnx-diagnostic 0.8.1__py3-none-any.whl → 0.8.3__py3-none-any.whl

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,7 @@ def onnx_dtype_to_torch_dtype(itype: int) -> "torch.dtype":  # noqa: F821
     )
 
 
-def torch_dtype_to_onnx_dtype(to: "torch.dtype") -> int:  # noqa: F821
+def torch_dtype_to_onnx_dtype(to: torch.dtype) -> int:
     """
     Converts a torch dtype into a onnx element type.
 
@@ -450,6 +449,11 @@ def fake_torchdynamo_exporting():
     """
     memorize = torch.compiler._is_exporting_flag
     torch.compiler._is_exporting_flag = True
+    assert torch.compiler.is_exporting(), (
+        f"Changes not detected "
+        f"torch.compiler._is_exporting_flag={torch.compiler._is_exporting_flag} "
+        f"and torch.compiler.is_exporting()={torch.compiler.is_exporting()}"
+    )
     try:
         yield
     finally:
@@ -478,7 +482,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()
@@ -493,6 +497,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
@@ -985,3 +1004,86 @@ 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