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

bias	Defines bias subclasses that work with scaled_dot_product_attention