less ambitious goal

Author: xadupre

.github/workflows/ci.yml

@@ -17,7 +17,7 @@ jobs:
         os: [ubuntu-latest]
         python: ['3.11', '3.12']
         transformers: ['4.48', '4.50', 'main']
-        torch: ['2.6', 'main']
+        torch: ['main']
 
     steps:
       - uses: actions/checkout@v3

_doc/examples/plot_export_locate_issue.py

@@ -25,6 +25,7 @@
 
 import traceback
 import torch
+from onnx_diagnostic import doc
 from onnx_diagnostic.torch_export_patches import bypass_export_some_errors
 
 
@@ -99,3 +100,6 @@ def forward(self, x: torch.Tensor, ys: list[torch.Tensor]):
 #
 #   File "onnx-diagnostic/_doc/examples/plot_export_locate_issue.py", line 25, in forward
 #       z = x * caty
+
+
+doc.plot_legend("was inferred to be a constant", "torch.export.export", "tomato")

_doc/examples/plot_export_tiny_llm.py

@@ -30,6 +30,7 @@
 import pprint
 import torch
 import transformers
+from onnx_diagnostic import doc
 from onnx_diagnostic.helpers import string_type
 from onnx_diagnostic.torch_models.llms import get_tiny_llm
 
@@ -170,3 +171,5 @@ def _forward_(*args, _f=None, **kwargs):
 # %%
 # If you have any error, then look at example
 # :ref:`l-plot-tiny-llm-export-patched`.
+
+doc.plot_legend("Tiny-LLM fails", "torch.export.export", "tomato")

_doc/examples/plot_export_tiny_llm_patched.py

@@ -66,6 +66,7 @@
 import pprint
 import torch
 import transformers
+from onnx_diagnostic import doc
 from onnx_diagnostic.helpers import string_type
 from onnx_diagnostic.torch_export_patches.onnx_export_errors import bypass_export_some_errors
 from onnx_diagnostic.torch_models.llms import get_tiny_llm
@@ -122,3 +123,6 @@
     )
     print("It worked:")
     print(ep)
+
+# %%
+doc.plot_legend("Tiny-LLM patched", "torch.export.export", "green")

_doc/examples/plot_export_with_dynamic_cache.py

@@ -23,6 +23,7 @@
 
 import pprint
 import torch
+from onnx_diagnostic import doc
 from onnx_diagnostic.cache_helpers import make_dynamic_cache
 from onnx_diagnostic.helpers import string_type
 from onnx_diagnostic.export import ModelInputs
@@ -221,3 +222,7 @@ def forward(self, cache, z):
             model, modificator(inputs[0]), dynamic_shapes=ds[0], strict=False
         )
 print(ep)
+
+# %%
+
+doc.plot_legend("dynamic shapes", "torch.export.export", "tomato")

_doc/examples/plot_export_with_dynamic_shapes_auto.py

@@ -12,6 +12,7 @@
 """
 
 import torch
+from onnx_diagnostic import doc
 
 
 class Model(torch.nn.Module):
@@ -90,3 +91,7 @@ def forward(self, x, y, z):
         dynamic_shapes=({0: AUTO, 1: AUTO}, {0: AUTO, 1: AUTO}, {0: AUTO, 1: AUTO}),
     )
 )
+
+# %%
+
+doc.plot_legend("dynamic shapes inferred", "torch.export.export", "tomato")

_doc/examples/plot_failing_model_extract.py

@@ -25,6 +25,7 @@
 import onnx
 import onnx.helper as oh
 import onnxruntime
+from onnx_diagnostic import doc
 from onnx_diagnostic.helpers import from_array_extended
 from onnx_diagnostic.ort_session import investigate_onnxruntime_issue
 
@@ -96,3 +97,7 @@
     onnx.shape_inference.infer_shapes(model, strict_mode=True)
 except onnx.onnx_cpp2py_export.shape_inference.InferenceError as e:
     print(e)
+
+# %%
+
+doc.plot_legend("Run until it fails", "onnxruntime.InferenceSession", "lightgrey")

_doc/examples/plot_failing_onnxruntime_evaluator.py

@@ -28,6 +28,7 @@
 import onnx.helper as oh
 import torch
 import onnxruntime
+from onnx_diagnostic import doc
 from onnx_diagnostic.ext_test_case import has_cuda
 from onnx_diagnostic.helpers import from_array_extended
 from onnx_diagnostic.reference import OnnxruntimeEvaluator
@@ -104,3 +105,5 @@
 # This runtime is useful when it fails for a numerical reason.
 # It is possible to insert prints in the python code to print
 # more information or debug if needed.
+
+doc.plot_legend("onnxruntime running step by step", "OnnxruntimeEvaluator", "lightgrey")

_doc/examples/plot_failing_reference_evaluator.py

@@ -19,6 +19,7 @@
 import onnx
 import onnx.helper as oh
 import onnxruntime
+from onnx_diagnostic import doc
 from onnx_diagnostic.helpers import from_array_extended
 from onnx_diagnostic.reference import ExtendedReferenceEvaluator
 
@@ -79,3 +80,5 @@
 # This runtime is useful when it fails for a numerical reason.
 # It is possible to insert prints in the python code to print
 # more information or debug if needed.
+
+doc.plot_legend("Python Runtime for ONNX", "ExtendedReferenceEvalutor", "lightgrey")