.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "auto_examples_backend/plot_custom_kernel.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note :ref:`Go to the end ` to download the full example code. .. rst-class:: sphx-glr-example-title .. _sphx_glr_auto_examples_backend_plot_custom_kernel.py: .. _l-example-plot-custom-kernel: Extend ReferenceEvaluator with a custom kernel ============================================== :class:`~onnx_light.onnx.reference.ReferenceEvaluator` dispatches every ``NodeProto`` against the static C++ ``KernelDispatchTable``. Any operator that is not built in — typically an operator from a user-defined domain, or a temporary stand-in for an op that is not yet implemented — would otherwise fail with ``unsupported op_type``. The :meth:`~onnx_light.onnx.reference.ReferenceEvaluator.register_custom_kernel` hook makes it possible to plug a Python callable into the runtime without touching the C++ dispatch table. The callable is invoked as ``fn(node, *numpy_inputs)`` and must return either a single :class:`numpy.ndarray` or a tuple/list of arrays for multi-output ops. Registrations are reapplied to the fresh ``RuntimeContext`` built on every :meth:`~onnx_light.onnx.reference.ReferenceEvaluator.run` call, so the same evaluator can be reused safely across runs. This example: * parses a small ONNX graph that calls ``my.domain.Scale`` — an operator that is *not* part of the built-in dispatch table, * registers a numpy-friendly Python implementation through :meth:`~onnx_light.onnx.reference.ReferenceEvaluator.register_custom_kernel`, * runs the model and prints the result. .. GENERATED FROM PYTHON SOURCE LINES 31-39 .. code-block:: Python from __future__ import annotations import numpy as np from onnx_light.onnx_lib import parser from onnx_light.onnx.reference import ReferenceEvaluator .. GENERATED FROM PYTHON SOURCE LINES 40-46 Build a model that uses a custom op +++++++++++++++++++++++++++++++++++ ``my.domain.Scale`` is a user-defined operator: it multiplies its single input by a ``factor`` attribute. The ONNX parser accepts the unknown op as long as the domain is declared in ``opset_import``. .. GENERATED FROM PYTHON SOURCE LINES 46-55 .. code-block:: Python model = parser.parse_model( '' "agraph (float[3] x) => (float[3] y) {" " y = my.domain.Scale(x)" "}" ) print(model) .. rst-class:: sphx-glr-script-out .. code-block:: none { ir_version: 10 opset_import: [ { version: 18 } { domain: "my.domain" version: 1 } ] graph: { name: "agraph" input: [ { name: "x" type: { tensor_type: { elem_type: 1 shape: { dim: [ { dim_value: 3 } ] } } } } ] output: [ { name: "y" type: { tensor_type: { elem_type: 1 shape: { dim: [ { dim_value: 3 } ] } } } } ] node: [ { input: ["x"] output: ["y"] op_type: "Scale" attribute: [ { factor: 3 } ] domain: "my.domain" } ] } } .. GENERATED FROM PYTHON SOURCE LINES 56-62 Implement the custom kernel +++++++++++++++++++++++++++ The callable receives the :class:`NodeProto` (so it can read attributes) followed by one ``numpy.ndarray`` per declared input. It returns either a single array or a tuple/list of arrays — one per declared output. .. GENERATED FROM PYTHON SOURCE LINES 62-69 .. code-block:: Python def scale(node, x): factor = next(a for a in node.attribute if str(a.name) == "factor").f return x * float(factor) .. GENERATED FROM PYTHON SOURCE LINES 70-78 Register and run ++++++++++++++++ :meth:`register_custom_kernel` is a wrapper around the low-level :py:meth:`RuntimeContext.register_custom_kernel` binding. The registration is stored on the evaluator and reapplied to every fresh :class:`RuntimeContext` it builds, so the same evaluator can run multiple inputs without re-registering. .. GENERATED FROM PYTHON SOURCE LINES 78-86 .. code-block:: Python sess = ReferenceEvaluator(model) sess.register_custom_kernel("my.domain", "Scale", scale) x = np.array([1.0, 2.0, 3.0], dtype=np.float32) (y,) = sess.run(None, {"x": x}) print(f"y = {y}") .. rst-class:: sphx-glr-script-out .. code-block:: none y = [3. 6. 9.] .. GENERATED FROM PYTHON SOURCE LINES 87-95 Overriding a built-in kernel ++++++++++++++++++++++++++++ A custom registration under the default ONNX domain takes precedence over the entry that ``KernelDispatchTable`` would otherwise dispatch. This is convenient to instrument or replace a specific kernel without patching the C++ runtime. Below ``Abs`` is replaced by negation just to demonstrate the override mechanism. .. GENERATED FROM PYTHON SOURCE LINES 95-111 .. code-block:: Python override_model = parser.parse_model( '' "agraph (float[3] x) => (float[3] y) { y = Abs(x) }" ) def fake_abs(node, x): return -x sess2 = ReferenceEvaluator(override_model) sess2.register_custom_kernel("", "Abs", fake_abs) (y2,) = sess2.run(None, {"x": np.array([-1.0, -2.0, -3.0], dtype=np.float32)}) print(f"Abs replaced by negation: y = {y2}") .. rst-class:: sphx-glr-script-out .. code-block:: none Abs replaced by negation: y = [1. 2. 3.] .. GENERATED FROM PYTHON SOURCE LINES 112-118 See also ++++++++ * :ref:`l-design-custom-kernels` for the design notes, including the matching low-level Python (``RuntimeContext.register_custom_kernel``) and C++ (``RuntimeContext::RegisterCustomKernel``) entry points. .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 0.004 seconds) .. _sphx_glr_download_auto_examples_backend_plot_custom_kernel.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: plot_custom_kernel.ipynb ` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: plot_custom_kernel.py ` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: plot_custom_kernel.zip ` .. include:: plot_custom_kernel.recommendations .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_