| .. 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:: and_masks |
| .. autofunction:: or_masks |
| .. autofunction:: noop_mask |
| |
| BlockMask |
| --------- |
| |
| .. autoclass:: BlockMask |
| :members: |
| :undoc-members: |
