Skip to content

nn.attention.flex_attention.md

Latest commit

 

History

History
45 lines (39 loc) · 705 Bytes

nn.attention.flex_attention.md

File metadata and controls

45 lines (39 loc) · 705 Bytes
.. 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

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