{eval-rst}

.. role:: hidden
    :class: hidden-section

torch.nn.attention.flex_attention

{eval-rst}

.. currentmodule:: torch.nn.attention.flex_attention

{eval-rst}

.. py:module:: torch.nn.attention.flex_attention

{eval-rst}

.. autofunction:: flex_attention

{eval-rst}

.. autoclass:: AuxOutput

{eval-rst}

.. autoclass:: AuxRequest

BlockMask Utilities

{eval-rst}

.. autofunction:: create_block_mask

{eval-rst}

.. autofunction:: create_mask

{eval-rst}

.. autofunction:: and_masks

{eval-rst}

.. autofunction:: or_masks

{eval-rst}

.. autofunction:: noop_mask

{eval-rst}

.. autoclass:: FlexKernelOptions
    :members:
    :undoc-members:

{eval-rst}

.. autoclass:: BlockMask
    :members:
    :undoc-members: