[fix] Add back BlockDiagonalMask import

It looks like BlockDiagonalMask import was accidentally removed in commit 166fd2e. This commit adds it back.
facebookresearch · Oct 22, 2024 · 86903f5 · 86903f5
1 parent 68b7fd1
commit 86903f5
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/xformers/ops/fmha/__init__.py b/xformers/ops/fmha/__init__.py
@@ -17,7 +17,12 @@
     flash3,
     triton_splitk,
 )
-from .attn_bias import VARLEN_BIASES, AttentionBias, LowerTriangularMask
+from .attn_bias import (
+    VARLEN_BIASES,
+    AttentionBias,
+    LowerTriangularMask,
+    BlockDiagonalMask,
+)
 from .common import (
     AttentionBwOpBase,
     AttentionFwOpBase,