torch.nn.attention¶
This module contains functions and classes that alter the behavior of torch.nn.functional.scaled_dot_product_attention
Utils¶
sdpa_kernel |
Context manager to select which backend to use for scaled dot product attention. |
SDPBackend |
An enum-like class that contains the different backends for scaled dot product attention. |
Submodules¶
This module implements the user facing API for flex_attention in PyTorch. |
|
Defines bias subclasses that work with scaled_dot_product_attention |
|
