documentation

Author: xadupre

_doc/conf.py

@@ -158,7 +158,7 @@
     "ignore_repr_types": "matplotlib\\.(text|axes)",
     # robubstness
     "reset_modules_order": "both",
-    "reset_modules": ("matplotlib", "onnx_diagnostic.reset_torch_transformers"),
+    "reset_modules": ("matplotlib", "onnx_diagnostic.doc.reset_torch_transformers"),
 }
 
 if int(os.environ.get("UNITTEST_GOING", "0")):

_doc/examples/plot_export_cond.py

@@ -13,6 +13,7 @@
 """
 
 import torch
+from onnx_diagnostic import doc
 
 
 # %%
@@ -84,3 +85,8 @@ def neg(x):
 
 ep = torch.export.export(model, (x,))
 print(ep.graph)
+
+
+# %%
+
+doc.plot_legend("If -> torch.cond", "torch.export.export", "tomato")

_doc/examples/plot_export_locate_issue.py

@@ -55,7 +55,8 @@ def forward(self, x: torch.Tensor, ys: list[torch.Tensor]):
 # %%
 # The error shows:
 #
-# ::
+# .. code-block::
+#
 #       Constraints violated (L['args'][0][0].size()[0])!
 #           For more information, run with TORCH_LOGS="+dynamic".
 #       - Not all values of RelaxedUnspecConstraint(L['args'][0][0].size()[0])
@@ -68,7 +69,7 @@ def forward(self, x: torch.Tensor, ys: list[torch.Tensor]):
 # One way to find the exact location is to retrieve a stack trace
 # by inserting an assert such as the following:
 #
-# ::
+# .. code-block::
 #
 #   assert msg != "range_refined_to_singleton", (
 #       f"A dynamic dimension becomes static! "
@@ -94,7 +95,7 @@ def forward(self, x: torch.Tensor, ys: list[torch.Tensor]):
 # is the following one. It points out the line turing a dynamic dimension into
 # static.
 #
-# ::
+# .. code-block::
 #
 #   File "onnx-diagnostic/_doc/examples/plot_export_locate_issue.py", line 25, in forward
 #       z = x * caty

_doc/examples/plot_export_tiny_llm.py

@@ -85,7 +85,8 @@ def _forward_(*args, _f=None, **kwargs):
 #
 # 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:`....get_tiny_llm`.
+# to create an untrained model:
+# :func:`onnx_diagnostic.torch_models.llms.get_tiny_llm`.
 # Then let's use it.
 
 experiment = get_tiny_llm()
@@ -139,7 +140,7 @@ def _forward_(*args, _f=None, **kwargs):
 #
 # Let's use the same dummy inputs but we use the downloaded model.
 # Dummy inputs and dynamic shapes are created by function
-# :func:`....get_tiny_llm`.
+# :func:`onnx_diagnostic.torch_models.llms.get_tiny_llm`.
 
 data = get_tiny_llm()
 inputs, dynamic_shapes = data["inputs"], data["dynamic_shapes"]

_doc/examples/plot_export_with_dynamic_shapes_auto.py

@@ -57,7 +57,7 @@ def forward(self, x, y, z):
         },
     )
     print(ep)
-    raise AssertionError("able to export this moel, please update the tutorial")
+    raise AssertionError("able to export this model, please update the tutorial")
 except torch._dynamo.exc.UserError as e:
     print(f"unable to use Dim('dz') because {type(e)}, {e}")
 

onnx_diagnostic/doc.py

@@ -0,0 +1,32 @@
+def reset_torch_transformers(gallery_conf, fname):
+    "Resets torch dynamo for :epkg:`sphinx-gallery`."
+    import matplotlib.pyplot as plt
+    import torch
+
+    plt.style.use("ggplot")
+    torch._dynamo.reset()
+
+
+def plot_legend(
+    text: str, text_bottom: str = "", color: str = "green", fontsize: int = 35
+) -> "matplotlib.axes.Axes":  # noqa: F821
+    import matplotlib.pyplot as plt
+
+    fig = plt.figure()
+    ax = fig.add_subplot()
+    ax.axis([0, 5, 0, 5])
+    ax.text(2.5, 4, "END", fontsize=50, horizontalalignment="center")
+    ax.text(
+        2.5,
+        2.5,
+        text,
+        fontsize=fontsize,
+        bbox={"facecolor": color, "alpha": 0.5, "pad": 10},
+        horizontalalignment="center",
+        verticalalignment="center",
+    )
+    if text_bottom:
+        ax.text(4.5, 0.5, text_bottom, fontsize=20, horizontalalignment="right")
+    ax.grid(False)
+    ax.set_axis_off()
+    return ax

pyproject.toml

@@ -14,6 +14,10 @@ exclude = [
     "^dist",  # skips dist
 ]
 
+[[tool.mypy.overrides]]
+module = ["onnx_diagnostic.doc"]
+disable_error_code = ["call-overload", "name-defined"]
+
 [[tool.mypy.overrides]]
 module = ["onnx_diagnostic.args"]
 disable_error_code = ["arg-type", "call-overload", "index"]