torch.nn.attention

此模块包含修改 torch.nn.functional.scaled_dot_product_attention 行为的函数和类

工具类

sdpa_kernel	用于选择用于扩展点积注意力的后端的上下文管理器。
SDPBackend	一个类似于枚举的类，包含扩展点积注意力的不同后端。

子模块

flex_attention	此模块实现了 PyTorch 中 flex_attention 的用户 API。
bias	定义与缩放点积注意力机制一起工作的偏置子类
experimental