doc

Author: xadupre

_doc/api/index.rst

@@ -8,6 +8,7 @@ API of onnx_diagnostic
     :caption: submodules
 
     torch_export_patches/index
+    torch_models/index
 
 .. toctree::
     :maxdepth: 1

_doc/api/torch_models/llms.rst

@@ -2,6 +2,6 @@
 onnx_diagnostic.torch_models.llms
 =================================
 
-.. automodule:: onnx_diagnostic.torch.models.llms
+.. automodule:: onnx_diagnostic.torch_models.llms
     :members:
     :no-undoc-members:

_doc/conf.py

@@ -113,6 +113,7 @@
     ("py:class", "torch.utils._pytree.Context"),
     ("py:class", "torch.utils._pytree.KeyEntry"),
     ("py:class", "torch.utils._pytree.TreeSpec"),
+    ("py:class", "transformers.LlamaConfig"),
     ("py:class", "transformers.cache_utils.Cache"),
     ("py:class", "transformers.cache_utils.DynamicCache"),
     ("py:class", "transformers.cache_utils.MambaCache"),

_doc/examples/plot_export_tiny_llm.py

@@ -1,4 +1,6 @@
 """
+.. _l-plot-tiny-llm-export:
+
 Export LLM with dynamic shapes
 ==============================
 

_unittests/ut_xrun_doc/test_documentation_examples.py

@@ -73,9 +73,9 @@ def add_test_methods(cls):
             if (
                 not reason
                 and name in {"plot_export_tiny_llm.py"}
-                and not has_transformers("4.50")
+                and not has_transformers("4.51")
             ):
-                reason = "transformers<4.50"
+                reason = "transformers<4.51"
 
             if reason:
 

onnx_diagnostic/torch_models/llms.py

@@ -18,6 +18,8 @@ def get_tiny_llm(
     :param dynamic_rope: use dynamic rope (see :class:`transformers.LlamaConfig`)
     :param kwargs: to overwrite the configuration, example ``num_hidden_layers=1``
     :return: dictionary
+
+    See :ref:`l-plot-tiny-llm-export` for an example.
     """
     config = {
         "architectures": ["LlamaForCausalLM"],