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

.. only:: html

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

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

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

.. _sphx_glr_examples_exp.py:


Exponential Function Example
========================

This example demonstrates how to implement an element-wise exponential function using Helion.

.. GENERATED FROM PYTHON SOURCE LINES 9-11

Imports
-------

.. GENERATED FROM PYTHON SOURCE LINES 11-20

.. code-block:: Python

    from __future__ import annotations

    import torch

    import helion
    from helion._testing import run_example
    import helion.language as hl



.. GENERATED FROM PYTHON SOURCE LINES 21-23

Exponential Kernel
---------------

.. GENERATED FROM PYTHON SOURCE LINES 23-40

.. code-block:: Python

    @helion.kernel()
    def exp(x: torch.Tensor) -> torch.Tensor:
        """
        Computes the exponential of all elements in the input tensor.

        Args:
            x: Input tensor

        Returns:
            Output tensor with the exponential of each element in the input
        """
        out = torch.empty_like(x)
        for tile in hl.tile(x.size()):
            out[tile] = torch.exp(x[tile])
        return out



.. GENERATED FROM PYTHON SOURCE LINES 41-43

Benchmark Wrapper
--------------

.. GENERATED FROM PYTHON SOURCE LINES 43-56

.. code-block:: Python

    def exp_tritonbench(x: torch.Tensor) -> dict[str, torch.Tensor]:
        """
        Wrapper for tritonbench that returns output in expected format.

        Args:
            x: Input tensor

        Returns:
            Dictionary containing the output tensor
        """
        return {"output": exp(x)}



.. GENERATED FROM PYTHON SOURCE LINES 57-59

Verification Function
-------------------

.. GENERATED FROM PYTHON SOURCE LINES 59-70

.. code-block:: Python

    def check(n: int) -> None:
        """
        Verify the exp kernel implementation against PyTorch's native exp function.

        Args:
            n: Size of the test tensor
        """
        x = torch.randn(n, device="cuda", dtype=torch.float32)
        run_example(exp, torch.exp, (x,))



.. GENERATED FROM PYTHON SOURCE LINES 71-73

Main Function
-----------

.. GENERATED FROM PYTHON SOURCE LINES 73-82

.. code-block:: Python

    def main() -> None:
        """
        Main entry point that runs the exp kernel verification with a tensor of size 1M elements.
        """
        check(1024 * 1024)


    if __name__ == "__main__":
        main()


.. _sphx_glr_download_examples_exp.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: exp.ipynb <exp.ipynb>`

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

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

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

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


.. only:: html

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

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