doc

Author: xadupre

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/api/onnx_tools.rst

@@ -0,0 +1,7 @@
+
+onnx_diagnostic.onnx_tools
+==========================
+
+.. automodule:: onnx_diagnostic.onnx_tools
+    :members:
+    :no-undoc-members:

_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
@@ -45,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/helpers.py

@@ -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