torch.nn.attention.flex_attention

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

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

.. autofunction:: flex_attention

BlockMask Utilities

.. autofunction:: create_block_mask

.. autofunction:: create_mask

.. autofunction:: and_masks

.. autofunction:: or_masks

.. autofunction:: noop_mask

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