better documentation

_doc/api/torch_export_patches/index.rst

@@ -1,6 +1,16 @@
 onnx_diagnostic.torch_export_patches
 ====================================
 
+.. toctree::
+    :maxdepth: 1
+    :caption: submodules
+
+    patches/index
+
 .. automodule:: onnx_diagnostic.torch_export_patches
     :members:
     :no-undoc-members:
+
+.. autofunction:: onnx_diagnostic.torch_export_patches.bypass_export_some_errors
+
+.. autofunction:: onnx_diagnostic.torch_export_patches.register_additional_serialization_functions

_doc/api/torch_export_patches/patches/index.rst

@@ -0,0 +1,13 @@
+onnx_diagnostic.torch_export_patches.patches
+============================================
+
+.. toctree::
+    :maxdepth: 1
+    :caption: modules
+
+    patch_torch
+    patch_transformers
+
+.. automodule:: onnx_diagnostic.torch_export_patches.patches
+    :members:
+    :no-undoc-members:

_doc/api/torch_export_patches/patches/patch_torch.rst

@@ -0,0 +1,7 @@
+
+onnx_diagnostic.torch_export_patches.patches.patch_torch
+========================================================
+
+.. automodule:: onnx_diagnostic.torch_export_patches.patches.patch_torch
+    :members:
+    :undoc-members:

_doc/api/torch_export_patches/patches/patch_transformers.rst

@@ -0,0 +1,6 @@
+onnx_diagnostic.torch_export_patches.patches.patch_transformers
+===============================================================
+
+.. automodule:: onnx_diagnostic.torch_export_patches.patches.patch_transformers
+    :members:
+    :undoc-members:

_doc/conf.py

@@ -86,6 +86,7 @@
     "numpy": ("https://numpy.org/doc/stable", None),
     "onnx": ("https://onnx.ai/onnx/", None),
     "onnx_array_api": ("https://sdpython.github.io/doc/onnx-array-api/dev/", None),
+    "onnx_diagnostic": ("https://sdpython.github.io/doc/onnx-diagnostic/dev/", None),
     "onnx_extended": ("https://sdpython.github.io/doc/onnx-extended/dev/", None),
     "pandas": ("https://pandas.pydata.org/pandas-docs/stable/", None),
     "python": (f"https://docs.python.org/{sys.version_info.major}", None),

_doc/examples/plot_export_tiny_llm.py

@@ -67,11 +67,10 @@ def _forward_(*args, _f=None, **kwargs):
 # The model creation
 # ++++++++++++++++++
 #
-# Let's create an untrained model.
-
-
-# %%
-# Let's get the model, inputs and dynamic shapes.
+# Let's create an untrained model using the config file provided
+# `config.json <https://huggingface.co/arnir0/Tiny-LLM/blob/main/config.json>`_
+# to create an untrained model: :func:`onnx_diagnostic.torch_models.llms.get_tiny_llm`.
+# Then let's use it.
 
 experiment = get_tiny_llm()
 untrained_model, inputs, dynamic_shapes = (
@@ -88,12 +87,15 @@ def _forward_(*args, _f=None, **kwargs):
 
 
 # %% Let's run it.
-print("input type", string_type(inputs, with_shape=True))
+print("input type before", string_type(inputs, with_shape=True))
 
 expected_output = untrained_model(**inputs)
 
+print("input type after-", string_type(inputs, with_shape=True))
+
+# %%
+# The outputs
 
-print("input after the execution", string_type(inputs, with_shape=True))
 print("result type", string_type(expected_output, with_shape=True))
 
 ep = torch.export.export(

onnx_diagnostic/__init__.py

@@ -3,5 +3,5 @@
 Functions, classes to dig into a model when this one is right, slow, wrong...
 """
 
-__version__ = "0.1.0"
+__version__ = "0.2.0"
 __author__ = "Xavier Dupré"