fix a few things

xadupre · xadupre · commit e6ac972f300f · 2025-03-21T11:35:12.000+01:00
diff --git a/.github/workflows/check-urls.yml b/.github/workflows/check-urls.yml
@@ -25,7 +25,7 @@ jobs:
     - name: urls-checker-code
       uses: urlstechie/urlchecker-action@master
       with:
-        subfolder: experimental_experiment
+        subfolder: onnx_diagnostic
         file_types: .md,.py,.rst,.ipynb
         print_all: false
         timeout: 2
diff --git a/_doc/api/args.rst b/_doc/api/args.rst
@@ -1,7 +1,7 @@
 
-.args
-=====
+onnx_diagnostic.args
+====================
 
-.. automodule:: experimental_experiment.args
+.. automodule:: onnx_diagnostic.args
     :members:
     :no-undoc-members:
diff --git a/_doc/api/cache_helpers.rst b/_doc/api/cache_helpers.rst
@@ -1,7 +1,7 @@
 
-.cache_helpers
-==============
+onnx_diagnostic.cache_helpers
+=============================
 
-.. automodule:: experimental_experiment.cache_helpers
+.. automodule:: onnx_diagnostic.cache_helpers
     :members:
     :no-undoc-members:
diff --git a/_doc/api/ext_test_case.rst b/_doc/api/ext_test_case.rst
@@ -1,7 +1,7 @@
 
-.ext_test_case
-==============
+onnx_diagnostic.ext_test_case
+=============================
 
-.. automodule:: experimental_experiment.ext_test_case
+.. automodule:: onnx_diagnostic.ext_test_case
     :members:
     :no-undoc-members:
diff --git a/_doc/api/helpers.rst b/_doc/api/helpers.rst
@@ -1,7 +1,7 @@
 
-.helpers
-========
+onnx_diagnostic.helpers
+=======================
 
-.. automodule:: experimental_experiment.helpers
+.. automodule:: onnx_diagnostic.helpers
     :members:
     :no-undoc-members:
diff --git a/_doc/api/index.rst b/_doc/api/index.rst
@@ -1,6 +1,6 @@
 
-API
-===
+API of onnx_diagnostic
+======================
 
 
 .. toctree::
@@ -17,6 +17,7 @@ API
     ext_test_case
     helpers
     ort_session
+    torch_test_helper
 
 
 .. automodule:: onnx_diagnostic
diff --git a/_doc/api/onnx_tools.rst b/_doc/api/onnx_tools.rst
@@ -1,7 +1,7 @@
 
-.onnx_tools
-===========
+onnx_diagnostic.onnx_tools
+==========================
 
-.. automodule:: experimental_experiment.onnx_tools
+.. automodule:: onnx_diagnostic.onnx_tools
     :members:
     :no-undoc-members:
diff --git a/_doc/api/ort_session.rst b/_doc/api/ort_session.rst
@@ -1,7 +1,7 @@
 
-.ort_session
-============
+onnx_diagnostic.ort_session
+===========================
 
-.. automodule:: experimental_experiment.ort_session
+.. automodule:: onnx_diagnostic.ort_session
     :members:
     :no-undoc-members:
diff --git a/_doc/api/torch_test_helper.rst b/_doc/api/torch_test_helper.rst
@@ -0,0 +1,7 @@
+
+onnx_diagnostic.torch_test_helper
+=================================
+
+.. automodule:: onnx_diagnostic.torch_test_helper
+    :members:
+    :no-undoc-members:
diff --git a/onnx_diagnostic/ext_test_case.py b/onnx_diagnostic/ext_test_case.py
@@ -228,7 +228,7 @@ def measure_time(
 
         from pprint import pprint
         from math import cos
-        from experimental_experiment.ext_test_case import measure_time
+        from onnx_diagnostic.ext_test_case import measure_time
 
         res = measure_time(lambda: cos(0.5))
         pprint(res)
@@ -349,13 +349,13 @@ def todo(cls, f: Callable, msg: str):
 
     def print_model(self, model: "ModelProto"):  # noqa: F821
         "Prints a ModelProto"
-        from experimental_experiment.helpers import pretty_onnx
+        from onnx_diagnostic.helpers import pretty_onnx
 
         print(pretty_onnx(model))
 
     def print_onnx(self, model: "ModelProto"):  # noqa: F821
         "Prints a ModelProto"
-        from experimental_experiment.helpers import pretty_onnx
+        from onnx_diagnostic.helpers import pretty_onnx
 
         print(pretty_onnx(model))
 
@@ -426,8 +426,8 @@ def assertEqualArrays(
 
     def assertEqualArray(
         self,
-        expected: numpy.ndarray,
-        value: numpy.ndarray,
+        expected: Any,
+        value: Any,
         atol: float = 0,
         rtol: float = 0,
         msg: Optional[str] = None,
@@ -1038,7 +1038,7 @@ def statistics_on_file(filename: str) -> Dict[str, Union[int, float, str]]:
         :showcode:
 
         import pprint
-        from experimental_experiment.ext_test_case import statistics_on_file, __file__
+        from onnx_diagnostic.ext_test_case import statistics_on_file, __file__
 
         pprint.pprint(statistics_on_file(__file__))
     """
@@ -1087,7 +1087,7 @@ def statistics_on_folder(
 
         import os
         import pprint
-        from experimental_experiment.ext_test_case import statistics_on_folder, __file__
+        from onnx_diagnostic.ext_test_case import statistics_on_folder, __file__
 
         pprint.pprint(statistics_on_folder(os.path.dirname(__file__)))
 
@@ -1099,7 +1099,7 @@ def statistics_on_folder(
 
         import os
         import pprint
-        from experimental_experiment.ext_test_case import statistics_on_folder, __file__
+        from onnx_diagnostic.ext_test_case import statistics_on_folder, __file__
 
         pprint.pprint(statistics_on_folder(os.path.dirname(__file__), aggregation=1))
     """
diff --git a/onnx_diagnostic/helpers.py b/onnx_diagnostic/helpers.py
@@ -7,6 +7,7 @@
 import numpy as np
 from onnx import (
     AttributeProto,
+    DataType,
     FunctionProto,
     GraphProto,
     ModelProto,
@@ -86,7 +87,7 @@ def size_type(dtype: Any) -> int:
     raise AssertionError(f"Unexpected dtype={dtype}")
 
 
-def tensor_dtype_to_np_dtype(tensor_dtype: int) -> np.dtype:
+def tensor_dtype_to_np_dtype(tensor_dtype: DataType) -> np.dtype:
     """
     Converts a TensorProto's data_type to corresponding numpy dtype.
     It can be used while making tensor.
@@ -140,7 +141,7 @@ def string_type(
     .. runpython::
         :showcode:
 
-        from experimental_experiment.helpers import string_type
+        from onnx_diagnostic.helpers import string_type
         print(string_type((1, ["r", 6.6])))
     """
     if obj is None:
@@ -402,9 +403,7 @@ def string_type(
 
 
 def string_signature(sig: Any) -> str:
-    """
-    Displays the signature of a functions.
-    """
+    """Displays the signature of a functions."""
 
     def _k(p, kind):
         for name in dir(p):
@@ -423,12 +422,10 @@ def _k(p, kind):
             t = f"*{t}"
         le = (30 - len(t)) * " "
         text.append(f"    {t}{le}|{_k(pp,kind)}")
-        text.append(
-            f") -> {sig.return_annotation}"
-            if sig.return_annotation is not inspect._empty
-            else ")"
-        )
-        return "\n".join(text)
+    text.append(
+        f") -> {sig.return_annotation}" if sig.return_annotation is not inspect._empty else ")"
+    )
+    return "\n".join(text)
 
 
 def string_sig(f: Callable, kwargs: Optional[Dict[str, Any]] = None) -> str:
@@ -585,7 +582,7 @@ def get_onnx_signature(model: ModelProto) -> Tuple[Tuple[str, Any], ...]:
     :param model: model
     :return: signature
     """
-    sig = []
+    sig: List[Any] = []
     for i in model.graph.input:
         dt = i.type
         if dt.HasField("sequence_type"):
@@ -782,7 +779,7 @@ def torch_dtype_to_onnx_dtype(to: "torch.dtype") -> int:  # noqa: F821
     raise NotImplementedError(f"Unable to convert torch dtype {to!r} to onnx dtype.")
 
 
-def dtype_to_tensor_dtype(dt: "dtype") -> int:  # noqa: F821
+def dtype_to_tensor_dtype(dt: Union[np.dtype, "torch.dtype"]) -> int:  # noqa: F821
     """
     Converts a torch dtype or numpy dtype into a onnx element type.
 
@@ -796,9 +793,9 @@ def dtype_to_tensor_dtype(dt: "dtype") -> int:  # noqa: F821
     return torch_dtype_to_onnx_dtype(dt)
 
 
-def np_dtype_to_tensor_dtype(dt: "dtype") -> int:  # noqa: F821
+def np_dtype_to_tensor_dtype(dt: np.dtype) -> int:  # noqa: F821
     """
-    Converts a tnumpy dtype into a onnx element type.
+    Converts a numpy dtype into a onnx element type.
 
     :param to: dtype
     :return: onnx type
@@ -848,8 +845,8 @@ def rename_dynamic_dimensions(
             common = v & original
             if not common:
                 continue
-            common = sorted(common)
-            by = common[0]
+            sorted_common = sorted(common)
+            by = sorted_common[0]
             if ban_prefix and by.startswith(ban_prefix):
                 continue
             replacements[k] = by
@@ -881,7 +878,7 @@ def visit_Name(self, node):
     return ast.unparse(new_tree)
 
 
-def flatten_object(x: Any, drop_keys: bool = False) -> List[Any]:
+def flatten_object(x: Any, drop_keys: bool = False) -> Any:
     """
     Flattens the object.
     It accepts some common classes used in deep learning.
@@ -914,7 +911,7 @@ def flatten_object(x: Any, drop_keys: bool = False) -> List[Any]:
         if isinstance(x.conv_states, list):
             res = flatten_object(x.conv_states) + flatten_object(x.ssm_states)
             return tuple(res)
-        return tuple(x.conv_states, x.ssm_states)
+        return (x.conv_states, x.ssm_states)
     if hasattr(x, "to_tuple"):
         return flatten_object(x.to_tuple(), drop_keys=drop_keys)
     if hasattr(x, "shape"):
diff --git a/onnx_diagnostic/onnx_tools.py b/onnx_diagnostic/onnx_tools.py
@@ -40,7 +40,7 @@ def onnx_lighten(
     """
     Creates a model without big initializers but stores statistics
     into dictionaries. The function can be reversed with
-    :func:`experimental_experiment.onnx_tools.onnx_unlighten`.
+    :func:`onnx_diagnostic.onnx_tools.onnx_unlighten`.
     The model is modified inplace.
 
     :param onx: model
@@ -97,7 +97,7 @@ def onnx_unlighten(
 ) -> ModelProto:
     """
     Function fixing the model produced by function
-    :func:`experimental_experiment.onnx_tools.onnx_lighten`.
+    :func:`onnx_diagnostic.onnx_tools.onnx_lighten`.
     The model is modified inplace.
 
     :param onx: model
diff --git a/onnx_diagnostic/ort_session.py b/onnx_diagnostic/ort_session.py
diff --git a/onnx_diagnostic/torch_test_helper.py b/onnx_diagnostic/torch_test_helper.py
diff --git a/pyproject.toml b/pyproject.toml
