.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "auto_examples/kernels/plot_custom_operators.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_auto_examples_kernels_plot_custom_operators.py>`
        to download the full example code.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_auto_examples_kernels_plot_custom_operators.py:


Run com.microsoft custom operators
==================================

This example registers the portable ``CDist``, ``BiasGelu``, and
``LinearAttention`` CPU kernels and runs one model containing the three
operators through onnx-light.

.. GENERATED FROM PYTHON SOURCE LINES 9-95




.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    Registered custom schemas: ['CDist']
    CDist output:
     [[1.        1.4142135 2.236068 ]
     [2.236068  3.1622777 1.       ]]
    BiasGelu output:
     [[-0.07010484 -0.10020959  0.84134424]
     [ 0.5800265   0.3457302   3.9998722 ]]
    LinearAttention output:
     [[[22.]
      [32.]]]






|

.. code-block:: Python


    # sphinx_gallery_thumbnail_path = "_static/gallery_thumbnails/custom_operators.png"

    import math

    import numpy as np

    from onnx_light.onnx import TensorProto, helper
    from onnx_light.onnx.reference import ReferenceEvaluator
    from onnx_light_cpu import operator_schema_lookup, register_kernels


    def make_model():
        """Builds a model containing the custom operators."""
        graph = helper.make_graph(
            [
                helper.make_node(
                    "CDist", ["A", "B"], ["distances"], domain="com.microsoft", metric="euclidean"
                ),
                helper.make_node("BiasGelu", ["X", "bias"], ["activated"], domain="com.microsoft"),
                helper.make_node(
                    "LinearAttention",
                    ["Q", "K", "V"],
                    ["attention", "state"],
                    domain="com.microsoft",
                    update_rule="linear",
                    q_num_heads=1,
                    kv_num_heads=1,
                    scale=1.0,
                ),
            ],
            "custom-operators",
            [
                helper.make_tensor_value_info("A", TensorProto.FLOAT, [None, None]),
                helper.make_tensor_value_info("B", TensorProto.FLOAT, [None, None]),
                helper.make_tensor_value_info("X", TensorProto.FLOAT, [None, None]),
                helper.make_tensor_value_info("bias", TensorProto.FLOAT, [None]),
                helper.make_tensor_value_info("Q", TensorProto.FLOAT, [1, 2, 2]),
                helper.make_tensor_value_info("K", TensorProto.FLOAT, [1, 2, 2]),
                helper.make_tensor_value_info("V", TensorProto.FLOAT, [1, 2, 1]),
            ],
            [
                helper.make_tensor_value_info("distances", TensorProto.FLOAT, [None, None]),
                helper.make_tensor_value_info("activated", TensorProto.FLOAT, [None, None]),
                helper.make_tensor_value_info("attention", TensorProto.FLOAT, [1, 2, 1]),
                helper.make_tensor_value_info("state", TensorProto.FLOAT, [1, 1, 2, 1]),
            ],
        )
        return helper.make_model(
            graph,
            opset_imports=[
                helper.make_opsetid("", 20),
                helper.make_opsetid("com.microsoft", 1),
            ],
            ir_version=13,
        )


    a = np.array([[0.0, 1.0], [2.0, 3.0]], dtype=np.float32)
    b = np.array([[1.0, 1.0], [-1.0, 2.0], [2.0, 2.0]], dtype=np.float32)
    x = np.array([[-2.0, -1.0, 0.0], [0.5, 1.0, 3.0]], dtype=np.float32)
    bias = np.array([0.25, -0.5, 1.0], dtype=np.float32)
    q = np.array([[[1.0, 2.0], [2.0, 1.0]]], dtype=np.float32)
    k = np.array([[[3.0, 4.0], [1.0, 2.0]]], dtype=np.float32)
    v = np.array([[[2.0], [3.0]]], dtype=np.float32)

    register_kernels()
    session = ReferenceEvaluator(make_model())
    distances, activated, attention, state = session.run(
        None, {"A": a, "B": b, "X": x, "bias": bias, "Q": q, "K": k, "V": v}
    )

    expected_distances = np.sqrt(np.sum((a[:, None, :] - b[None, :, :]) ** 2, axis=2))
    z = x + bias
    expected_activated = (
        0.5 * z * (1.0 + np.vectorize(math.erf, otypes=[np.float32])(z / np.sqrt(np.float32(2.0))))
    )
    np.testing.assert_allclose(distances, expected_distances, rtol=1e-6, atol=1e-6)
    np.testing.assert_allclose(activated, expected_activated, rtol=1e-6, atol=1e-5)
    np.testing.assert_allclose(attention, np.array([[[22.0], [32.0]]], dtype=np.float32))
    np.testing.assert_allclose(state, np.array([[[[9.0], [14.0]]]], dtype=np.float32))

    print("Registered custom schemas:", [schema.name for schema in operator_schema_lookup("CDist")])
    print("CDist output:\n", distances)
    print("BiasGelu output:\n", activated)
    print("LinearAttention output:\n", attention)


.. rst-class:: sphx-glr-timing

   **Total running time of the script:** (0 minutes 0.055 seconds)


.. _sphx_glr_download_auto_examples_kernels_plot_custom_operators.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_operators.ipynb <plot_custom_operators.ipynb>`

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: plot_custom_operators.py <plot_custom_operators.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: plot_custom_operators.zip <plot_custom_operators.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_