First PR

.github/workflows/documentation.yml

@@ -82,7 +82,6 @@ jobs:
         run: cat doc.txt
 
       - name: Check for errors and warnings
-        continue-on-error: true
         run: |
           if [[ $(grep ERROR doc.txt | grep -v 'Unknown target name: "l_shape"' | grep -v 'Unknown target name: "l_x"') ]]; then
             echo "Documentation produces errors."

.github/workflows/mypy.yml

@@ -0,0 +1,14 @@
+name: Type annotation with mypy
+on: [push, pull_request]
+jobs:
+  mypy:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - uses: actions/setup-python@v4
+        with:
+          python-version: '3.12'
+      - name: Install mypy
+        run: pip install mypy
+      - name: Run mypy
+        run: mypy

README.rst

@@ -44,6 +44,51 @@ or
 
     pip install onnx-diagnostic
 
+Snapshot of usefuls tools
++++++++++++++++++++++++++
+
+**string_type**
+
+.. code-block:: python
+
+    import torch
+    from onnx_diagnostic.helpers import string_type
+
+    inputs = (
+        torch.rand((3, 4), dtype=torch.float16),
+        [
+            torch.rand((5, 6), dtype=torch.float16),
+            torch.rand((5, 6, 7), dtype=torch.float16),
+        ]
+    )
+
+    # with shapes
+    print(string_type(inputs, with_shape=True))
+
+::
+
+    >>> (T10s3x4,#2[T10s5x6,T10s5x6x7])
+
+**onnx_dtype_name**
+
+.. code-block:: python
+
+        import onnx
+        from onnx_diagnostic.helpers import onnx_dtype_name
+
+        itype = onnx.TensorProto.BFLOAT16
+        print(onnx_dtype_name(itype))
+        print(onnx_dtype_name(7))
+
+::
+
+    >>> BFLOAT16
+    >>> INT64
+
+**max_diff**
+
+Returns the maximum discrancies accross nested containers containing tensors.
+
 Documentation
 +++++++++++++
 

_doc/api/index.rst

@@ -16,6 +16,7 @@ API of onnx_diagnostic
     cache_helpers
     ext_test_case
     helpers
+    onnx_tools
     ort_session
     torch_test_helper
 

_doc/conf.py

@@ -85,10 +85,6 @@
     "matplotlib": ("https://matplotlib.org/stable/", None),
     "numpy": ("https://numpy.org/doc/stable", None),
     "onnx": ("https://onnx.ai/onnx/", None),
-    "onnx_diagnostic": (
-        "https://sdpython.github.io/doc/onnx-diagnostic/dev/",
-        None,
-    ),
     "onnx_array_api": ("https://sdpython.github.io/doc/onnx-array-api/dev/", None),
     "onnx_extended": ("https://sdpython.github.io/doc/onnx-extended/dev/", None),
     "pandas": ("https://pandas.pydata.org/pandas-docs/stable/", None),
@@ -118,8 +114,8 @@
     ("py:class", "torch.utils._pytree.KeyEntry"),
     ("py:class", "torch.utils._pytree.TreeSpec"),
     ("py:class", "transformers.cache_utils.Cache"),
-    # ("py:class", "transformers.cache_utils.DynamicCache"),
-    # ("py:class", "transformers.cache_utils.MambaCache"),
+    ("py:class", "transformers.cache_utils.DynamicCache"),
+    ("py:class", "transformers.cache_utils.MambaCache"),
     ("py:func", "torch.export._draft_export.draft_export"),
     ("py:func", "torch._export.tools.report_exportability"),
 ]
@@ -128,7 +124,9 @@
     ("py:func", ".*numpy[.].*"),
     ("py:func", ".*scipy[.].*"),
     # ("py:func", ".*torch.ops.higher_order.*"),
+    ("py:class", ".*numpy._typing[.].*"),
     ("py:class", ".*onnxruntime[.].*"),
+    ("py:meth", ".*onnxruntime[.].*"),
 ]
 
 
@@ -148,7 +146,7 @@
     # errors
     "abort_on_example_error": True,
     # recommendation
-    "recommender": {"enable": True, "n_examples": 5, "min_df": 3, "max_df": 0.9},
+    "recommender": {"enable": True, "n_examples": 3, "min_df": 3, "max_df": 0.9},
     # ignore capture for matplotib axes
     "ignore_repr_types": "matplotlib\\.(text|axes)",
     # robubstness

_doc/examples/plot_exporter_exporter_dynamic_shapes_auto.py

@@ -1,6 +1,4 @@
 """
-.. _l-plot-exporter-dynamic_shapes:
-
 Use DYNAMIC or AUTO when dynamic shapes has constraints
 =======================================================
 

_doc/index.rst

@@ -1,6 +1,6 @@
 
-onnx-diagnostic: fuzzy work
-===================================
+onnx-diagnostic: investigate onnx models
+========================================
 
 .. image:: https://github.com/sdpython/onnx-diagnostic/actions/workflows/documentation.yml/badge.svg
     :target: https://github.com/sdpython/onnx-diagnostic/actions/workflows/documentation.yml
@@ -36,6 +36,7 @@ Source are `sdpython/onnx-diagnostic
     :caption: Contents
 
     api/index
+    galleries
 
 .. toctree::
     :maxdepth: 1
@@ -44,6 +45,45 @@ Source are `sdpython/onnx-diagnostic
     CHANGELOGS
     license
 
+
+**Some usefuls tools**
+
+.. code-block:: python
+
+    import torch
+    from onnx_diagnostic.helpers import string_type
+
+    inputs = (
+        torch.rand((3, 4), dtype=torch.float16),
+        [
+            torch.rand((5, 6), dtype=torch.float16),
+            torch.rand((5, 6, 7), dtype=torch.float16),
+        ]
+    )
+
+    # with shapes
+    print(string_type(inputs, with_shape=True))
+
+::
+
+    >>> (T10s3x4,#2[T10s5x6,T10s5x6x7])
+
+.. code-block:: python
+
+        import onnx
+        from onnx_diagnostic.helpers import onnx_dtype_name
+
+        itype = onnx.TensorProto.BFLOAT16
+        print(onnx_dtype_name(itype))
+        print(onnx_dtype_name(7))
+
+::
+
+    >>> BFLOAT16
+    >>> INT64
+
+:func:`onnx_diagnostic.helpers.max_diff`, ...
+
 The documentation was updated on:
 
 .. runpython::

onnx_diagnostic/ext_test_case.py

@@ -4,7 +4,7 @@
 """
 
 import glob
-import importlib
+import importlib.util
 import logging
 import os
 import re

onnx_diagnostic/helpers.py

@@ -5,9 +5,9 @@
 import sys
 from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
 import numpy as np
+import numpy.typing as npt
 from onnx import (
     AttributeProto,
-    DataType,
     FunctionProto,
     GraphProto,
     ModelProto,
@@ -87,7 +87,7 @@ def size_type(dtype: Any) -> int:
     raise AssertionError(f"Unexpected dtype={dtype}")
 
 
-def tensor_dtype_to_np_dtype(tensor_dtype: DataType) -> np.dtype:
+def tensor_dtype_to_np_dtype(tensor_dtype: int) -> np.dtype:
     """
     Converts a TensorProto's data_type to corresponding numpy dtype.
     It can be used while making tensor.
@@ -105,7 +105,7 @@ def tensor_dtype_to_np_dtype(tensor_dtype: DataType) -> np.dtype:
                 f"ml_dtypes can be used."
             ) from e
 
-        mapping = {
+        mapping: Dict[int, np.dtype] = {
             TensorProto.BFLOAT16: ml_dtypes.bfloat16,
             TensorProto.FLOAT8E4M3FN: ml_dtypes.float8_e4m3fn,
             TensorProto.FLOAT8E4M3FNUZ: ml_dtypes.float8_e4m3fnuz,
@@ -142,7 +142,30 @@ def string_type(
         :showcode:
 
         from onnx_diagnostic.helpers import string_type
+
         print(string_type((1, ["r", 6.6])))
+
+    With pytorch:
+
+    .. runpython::
+        :showcode:
+
+        import torch
+        from onnx_diagnostic.helpers import string_type
+
+        inputs = (
+            torch.rand((3, 4), dtype=torch.float16),
+            [
+                torch.rand((5, 6), dtype=torch.float16),
+                torch.rand((5, 6, 7), dtype=torch.float16),
+            ]
+        )
+
+        # with shapes
+        print(string_type(inputs, with_shape=True))
+
+        # with min max
+        print(string_type(inputs, with_shape=True, with_min_max=True))
     """
     if obj is None:
         return "None"
@@ -465,7 +488,19 @@ def string_sig(f: Callable, kwargs: Optional[Dict[str, Any]] = None) -> str:
 
 @functools.cache
 def onnx_dtype_name(itype: int) -> str:
-    """Returns the ONNX name for a specific element type."""
+    """
+    Returns the ONNX name for a specific element type.
+
+    .. runpython::
+        :showcode:
+
+        import onnx
+        from onnx_diagnostic.helpers import onnx_dtype_name
+
+        itype = onnx.TensorProto.BFLOAT16
+        print(onnx_dtype_name(itype))
+        print(onnx_dtype_name(7))
+    """
     for k in dir(TensorProto):
         v = getattr(TensorProto, k)
         if v == itype:
@@ -477,19 +512,24 @@ def pretty_onnx(
     onx: Union[FunctionProto, GraphProto, ModelProto, ValueInfoProto, str],
     with_attributes: bool = False,
     highlight: Optional[Set[str]] = None,
+    shape_inference: bool = False,
 ) -> str:
     """
     Displays an onnx prot in a better way.
 
     :param with_attributes: displays attributes as well, if only a node is printed
     :param highlight: to highlight some names
+    :param shape_inference: run shape inference before printing the model
     :return: text
     """
     assert onx is not None, "onx cannot be None"
     if isinstance(onx, str):
         onx = onnx_load(onx, load_external_data=False)
     assert onx is not None, "onx cannot be None"
 
+    if shape_inference:
+        onx = onx.shape_inference.infer_shapes(onx)
+
     if isinstance(onx, ValueInfoProto):
         name = onx.name
         itype = onx.type.tensor_type.elem_type
@@ -577,7 +617,7 @@ def make_hash(obj: Any) -> str:
 
 def get_onnx_signature(model: ModelProto) -> Tuple[Tuple[str, Any], ...]:
     """
-    Produces a tuple of tuples correspinding to the signatures.
+    Produces a tuple of tuples corresponding to the signatures.
 
     :param model: model
     :return: signature
@@ -611,7 +651,7 @@ def convert_endian(tensor: TensorProto) -> None:
     tensor.raw_data = np.frombuffer(tensor.raw_data, dtype=np_dtype).byteswap().tobytes()
 
 
-def from_array_ml_dtypes(arr: np.ndarray, name: Optional[str] = None) -> TensorProto:
+def from_array_ml_dtypes(arr: npt.ArrayLike, name: Optional[str] = None) -> TensorProto:
     """
     Converts a numpy array to a tensor def assuming the dtype
     is defined in ml_dtypes.
@@ -625,7 +665,7 @@ def from_array_ml_dtypes(arr: np.ndarray, name: Optional[str] = None) -> TensorP
     """
     import ml_dtypes
 
-    assert isinstance(arr, np.ndarray), f"arr must be of type np.ndarray, got {type(arr)}"
+    assert isinstance(arr, np.ndarray), f"arr must be of type numpy.ndarray, got {type(arr)}"
 
     tensor = TensorProto()
     tensor.dims.extend(arr.shape)
@@ -651,9 +691,9 @@ def from_array_ml_dtypes(arr: np.ndarray, name: Optional[str] = None) -> TensorP
     return tensor
 
 
-def from_array_extended(tensor: np.ndarray, name: Optional[str] = None) -> TensorProto:
+def from_array_extended(tensor: npt.ArrayLike, name: Optional[str] = None) -> TensorProto:
     """
-    Converts an array into a TensorProto.
+    Converts an array into a :class:`onnx.TensorProto`.
 
     :param tensor: numpy array
     :param name: name

onnx_diagnostic/ort_session.py

@@ -1,6 +1,7 @@
 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
@@ -128,7 +129,7 @@ def __init__(
 class InferenceSessionForNumpy(_InferenceSession):
     """
     Wraps an `onnxruntime.InferenceSession` to overload method `run`
-    to support :class:`np.ndarray`.
+    to support :class:`numpy.ndarray`.
 
     :param sess: model or inference session
     :param session_options: options
@@ -172,8 +173,8 @@ def __init__(
         )
 
     def run(
-        self, output_names: Optional[List[str]], feeds: Dict[str, np.ndarray]
-    ) -> List[np.ndarray]:
+        self, output_names: Optional[List[str]], feeds: Dict[str, npt.ArrayLike]
+    ) -> List[npt.ArrayLike]:
         """Calls :meth:`onnxruntime.InferenceSession.run`."""
         return self.sess.run(output_names, feeds)
 
@@ -293,7 +294,7 @@ def _get_ortvalues_from_torch_tensors(
 
     def _ortvalues_to_torch_tensor(
         self,
-        ortvalues: Union[List[onnxruntime.OrtValue], onnxruntime.OrtValueVector],
+        ortvalues: Union[List[ORTC.OrtValue], ORTC.OrtValueVector],
     ) -> Tuple[torch.Tensor, ...]:
         if len(ortvalues) == 0:
             return tuple()
@@ -403,7 +404,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, np.ndarray]]] = None,
+    feeds: Optional[Union[Dict[str, torch.Tensor], Dict[str, npt.ArrayLike]]] = None,
     verbose: int = 0,
     dump_filename: Optional[str] = None,
     infer_shapes: bool = True,