.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "auto_examples_runtime/plot_register_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_runtime_plot_register_custom_kernel.py: .. _l-example-plot-register-custom-kernel: Replace a built-in kernel with a Python one and prove it ran ============================================================ This is the Python counterpart of the standalone C++ example ``examples/register_custom_kernel`` (see :ref:`l-cpp-register-custom-kernel-example`). The C++ example writes a brand-new :cpp:class:`~onnx_light::core::runtime::KernelBase` subclass for the existing ``Abs`` operator, installs it into the shared dispatch table and checks — via a run counter — that the custom kernel, not the built-in one, executed the node. Here the same scenario is expressed in Python: * a one-node ``y = Abs(x)`` model is parsed, * a numpy-friendly implementation of ``Abs`` is registered through :meth:`~onnx_light.onnx.reference.ReferenceEvaluator.register_custom_kernel` under the default ONNX domain, overriding the built-in kernel, * the model runs and the example asserts both that ``y == |x|`` and that the custom kernel — bumping a shared run counter on every call — was the one dispatched. It also prints the custom kernel's ``":::"`` identifier next to the official built-in one and checks the two differ. .. GENERATED FROM PYTHON SOURCE LINES 27-35 .. 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 36-41 Build a one-node ``Abs`` model ++++++++++++++++++++++++++++++ ``Abs`` is a built-in operator, so without an override the built-in kernel would compute the result. .. GENERATED FROM PYTHON SOURCE LINES 41-48 .. code-block:: Python model = parser.parse_model( '' "agraph (float[4] x) => (float[4] y) { y = Abs(x) }" ) print(model) .. rst-class:: sphx-glr-script-out .. code-block:: none { ir_version: 10 opset_import: [ { domain: "" version: 18 } ] graph: { name: "agraph" input: [ { name: "x" type: { tensor_type: { elem_type: 1 shape: { dim: [ { dim_value: 4 } ] } } } } ] output: [ { name: "y" type: { tensor_type: { elem_type: 1 shape: { dim: [ { dim_value: 4 } ] } } } } ] node: [ { input: [ "x" ] output: [ "y" ] op_type: "Abs" domain: "" } ] } } .. GENERATED FROM PYTHON SOURCE LINES 49-56 Implement the custom kernel +++++++++++++++++++++++++++ The callable receives the :class:`NodeProto` followed by one ``numpy.ndarray`` per input and returns the element-wise absolute value. A module-level counter records every invocation so the example can prove the custom kernel — and not the built-in ``Abs`` — actually ran. .. GENERATED FROM PYTHON SOURCE LINES 56-77 .. code-block:: Python run_count = 0 # Built-in kernels expose a ``":::"`` # identifier (see the C++ classes, e.g. the official ``Abs`` kernel is named # ``"onnx_kernels:CPU:ai.onnx:Abs"``). The custom kernel below advertises its # own name under a distinct ``example`` library prefix so it never collides # with — and is clearly distinguishable from — the built-in one. OFFICIAL_ABS_KERNEL_NAME = "onnx_kernels:CPU:ai.onnx:Abs" CUSTOM_ABS_KERNEL_NAME = "example:CPU:ai.onnx:Abs" def custom_abs(node, x): global run_count run_count += 1 return np.abs(x) custom_abs.name = CUSTOM_ABS_KERNEL_NAME .. GENERATED FROM PYTHON SOURCE LINES 78-84 Show the custom kernel name differs from the official one +++++++++++++++++++++++++++++++++++++++++++++++++++++++++ Both follow the same ``":::"`` convention but use different library prefixes, so the override is unmistakably *not* the built-in ``Abs`` kernel. .. GENERATED FROM PYTHON SOURCE LINES 84-93 .. code-block:: Python print(f"official Abs kernel name: {OFFICIAL_ABS_KERNEL_NAME}") print(f"custom Abs kernel name: {custom_abs.name}") assert custom_abs.name != OFFICIAL_ABS_KERNEL_NAME, ( "the custom kernel name must differ from the official one, otherwise it " "would be indistinguishable from the built-in kernel." ) print("OK: the custom kernel name is different from the official one.") .. rst-class:: sphx-glr-script-out .. code-block:: none official Abs kernel name: onnx_kernels:CPU:ai.onnx:Abs custom Abs kernel name: example:CPU:ai.onnx:Abs OK: the custom kernel name is different from the official one. .. GENERATED FROM PYTHON SOURCE LINES 94-99 Register the override and run +++++++++++++++++++++++++++++ Registering under the default ONNX domain (the empty string is normalised to ``ai.onnx``) takes precedence over the built-in ``Abs`` entry. .. GENERATED FROM PYTHON SOURCE LINES 99-107 .. code-block:: Python sess = ReferenceEvaluator(model) sess.register_custom_kernel("", "Abs", custom_abs) x = np.array([-1.0, 2.0, -3.5, 0.0], dtype=np.float32) (y,) = sess.run(None, {"x": x}) print(f"y = {y}") .. rst-class:: sphx-glr-script-out .. code-block:: none y = [1. 2. 3.5 0. ] .. GENERATED FROM PYTHON SOURCE LINES 108-113 Verify the output and that the custom kernel ran ++++++++++++++++++++++++++++++++++++++++++++++++ The output must equal ``|x|`` and the run counter must show the custom kernel was dispatched exactly once, proving the override replaced the built-in. .. GENERATED FROM PYTHON SOURCE LINES 113-121 .. code-block:: Python np.testing.assert_allclose(y, np.abs(x)) assert run_count == 1, ( f"expected the custom Abs kernel to run exactly once, but it ran " f"{run_count} time(s); the built-in kernel was probably dispatched instead." ) print(f"PASS: the custom 'Abs' kernel ran {run_count} time(s) and produced |x|.") .. rst-class:: sphx-glr-script-out .. code-block:: none PASS: the custom 'Abs' kernel ran 1 time(s) and produced |x|. .. GENERATED FROM PYTHON SOURCE LINES 122-129 See also ++++++++ * :ref:`l-example-plot-custom-kernel` for a broader tour of custom and user-defined-domain kernels in Python. * :ref:`l-cpp-register-custom-kernel-example` for the equivalent standalone C++ example. .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 0.003 seconds) .. _sphx_glr_download_auto_examples_runtime_plot_register_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_register_custom_kernel.ipynb ` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: plot_register_custom_kernel.py ` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: plot_register_custom_kernel.zip ` .. include:: plot_register_custom_kernel.recommendations .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_