torch.nn.attention

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

实用程序

sdpa_kernel	Context manager 来选择用于缩放点积注意力的后端。
SDP下载结束	一个类似枚举的类，包含用于缩放点积注意力的不同后端。

子模块

flex_attention	此模块在 PyTorch 中实现面向用户的 API for flex_attention。
偏见	定义与 scaled_dot_product_attention 一起使用的 bias 子类