nn.attention.flex_attention.md

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

torch.nn.attention.flex_attention

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

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

.. autofunction:: flex_attention

.. autoclass:: AuxOutput

.. autoclass:: AuxRequest

BlockMask Utilities

.. autofunction:: create_block_mask

.. autofunction:: create_mask

.. autofunction:: and_masks

.. autofunction:: or_masks

.. autofunction:: noop_mask

FlexKernelOptions

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

BlockMask

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