docuemntation

Author: xadupre

_doc/api/reference/torch_ops/control_flow_ops.rst

@@ -0,0 +1,6 @@
+
+onnx_diagnostic.reference.torch_ops.controlflow_ops
+===================================================
+
+.. automodule:: onnx_diagnostic.reference.torch_ops.controlflow_ops
+    :members:

_doc/api/reference/torch_ops/controlflow_ops.rst

@@ -9,7 +9,7 @@ onnx_diagnostic.reference.torch_ops
 
     access_ops
     binary_ops
-    control_flow_ops
+    controlflow_ops
     generator_ops
     nn_ops
     other_ops

_doc/api/reference/torch_ops/index.rst

@@ -2,6 +2,7 @@
 import unittest
 from onnx_diagnostic.ext_test_case import (
     ExtTestCase,
+    ignore_errors,
     requires_torch,
     requires_transformers,
     hide_stdout,
@@ -22,6 +23,7 @@ class TestModelBuilderHelper(ExtTestCase):
     # This is to limit impact on CI.
     @requires_transformers("4.52")
     @requires_torch("2.7.99")
+    @ignore_errors(OSError)  # connectivity issues
     def test_download_model_builder(self):
         path = download_model_builder_to_cache()
         self.assertExists(path)
@@ -32,6 +34,7 @@ def test_download_model_builder(self):
     @requires_transformers("4.52")
     @requires_torch("2.7.99")
     @hide_stdout()
+    @ignore_errors(OSError)  # connectivity issues
     def test_model_builder_id(self):
         # clear&&python ~/.cache/onnx-diagnostic/builder.py
         # --model arnir0/Tiny-LLM -p fp16 -c dump_cache -e cpu -o dump_model

_unittests/ut_helpers/test_model_builder_helper.py

@@ -0,0 +1,30 @@
+import onnx
+import onnx.helper as oh
+import torch
+from onnx_diagnostic.helpers import string_type
+from onnx_diagnostic.reference import TorchOnnxEvaluator
+
+TFLOAT = onnx.TensorProto.FLOAT
+
+proto = oh.make_model(
+    oh.make_graph(
+        [
+            oh.make_node("Sigmoid", ["Y"], ["sy"]),
+            oh.make_node("Mul", ["Y", "sy"], ["ysy"]),
+            oh.make_node("Mul", ["X", "ysy"], ["final"]),
+        ],
+        "-nd-",
+        [
+            oh.make_tensor_value_info("X", TFLOAT, [1, "b", "c"]),
+            oh.make_tensor_value_info("Y", TFLOAT, ["a", "b", "c"]),
+        ],
+        [oh.make_tensor_value_info("final", TFLOAT, ["a", "b", "c"])],
+    ),
+    opset_imports=[oh.make_opsetid("", 18)],
+    ir_version=9,
+)
+
+sess = TorchOnnxEvaluator(proto, verbose=1)
+feeds = dict(X=torch.rand((4, 5)), Y=torch.rand((4, 5)))
+result = sess.run(None, feeds)
+print(string_type(result, with_shape=True, with_min_max=True))

k.py

@@ -62,6 +62,87 @@ class TorchOnnxEvaluator:
     The class is not multithreaded. `runtime_info` gets updated
     by the the class. The list of available kernels is returned by function
     :func:`onnx_diagnostic.reference.torch_evaluator.get_kernels`.
+    Example:
+
+    .. runpython::
+        :showcode:
+
+        import onnx
+        import onnx.helper as oh
+        import torch
+        from onnx_diagnostic.helpers import string_type
+        from onnx_diagnostic.reference import TorchOnnxEvaluator
+
+        TFLOAT = onnx.TensorProto.FLOAT
+
+        proto = oh.make_model(
+            oh.make_graph(
+                [
+                    oh.make_node("Sigmoid", ["Y"], ["sy"]),
+                    oh.make_node("Mul", ["Y", "sy"], ["ysy"]),
+                    oh.make_node("Mul", ["X", "ysy"], ["final"]),
+                ],
+                "-nd-",
+                [
+                    oh.make_tensor_value_info("X", TFLOAT, [1, "b", "c"]),
+                    oh.make_tensor_value_info("Y", TFLOAT, ["a", "b", "c"]),
+                ],
+                [oh.make_tensor_value_info("final", TFLOAT, ["a", "b", "c"])],
+            ),
+            opset_imports=[oh.make_opsetid("", 18)],
+            ir_version=9,
+        )
+
+        sess = TorchOnnxEvaluator(proto)
+        feeds = dict(X=torch.rand((4, 5)), Y=torch.rand((4, 5)))
+        result = sess.run(None, feeds)
+        print(string_type(result, with_shape=True, with_min_max=True))
+
+    Adding ``verbose=1`` shows which kernels is executed:
+
+    .. runpython::
+        :showcode:
+
+        import onnx
+        import onnx.helper as oh
+        import torch
+        from onnx_diagnostic.helpers import string_type
+        from onnx_diagnostic.reference import TorchOnnxEvaluator
+
+        TFLOAT = onnx.TensorProto.FLOAT
+
+        proto = oh.make_model(
+            oh.make_graph(
+                [
+                    oh.make_node("Sigmoid", ["Y"], ["sy"]),
+                    oh.make_node("Mul", ["Y", "sy"], ["ysy"]),
+                    oh.make_node("Mul", ["X", "ysy"], ["final"]),
+                ],
+                "-nd-",
+                [
+                    oh.make_tensor_value_info("X", TFLOAT, [1, "b", "c"]),
+                    oh.make_tensor_value_info("Y", TFLOAT, ["a", "b", "c"]),
+                ],
+                [oh.make_tensor_value_info("final", TFLOAT, ["a", "b", "c"])],
+            ),
+            opset_imports=[oh.make_opsetid("", 18)],
+            ir_version=9,
+        )
+
+        sess = TorchOnnxEvaluator(proto, verbose=1)
+        feeds = dict(X=torch.rand((4, 5)), Y=torch.rand((4, 5)))
+        result = sess.run(None, feeds)
+        print(string_type(result, with_shape=True, with_min_max=True))
+
+    It also shows when a result is not needed anymore. In that case,
+    it is deleted to free the memory it takes.
+    The runtime can also execute the kernel the onnx model on CUDA.
+    It follows the same logic as :class:`onnxruntime.InferenceSession`:
+    ``providers=["CUDAExecutionProvider"]``.
+    It is better in that case to move the input on CUDA. The class
+    tries to move every weight on CUDA but tries to keep any tensor
+    identified as a shape in CPU. Some bugs may remain as torch
+    raises an exception when devices are expected to be the same.
     """
 
     class IO: