Shortcuts

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

flex_attention

This module implements the user facing API for flex_attention in PyTorch.

bias

Defines bias subclasses that work with scaled_dot_product_attention

Docs

Access comprehensive developer documentation for PyTorch

View Docs

Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

Resources

Find development resources and get your questions answered

View Resources