nn.attention.flex_attention.rst

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:: create_nested_block_mask

.. autofunction:: and_masks

.. autofunction:: or_masks

.. autofunction:: noop_mask

BlockMask

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