make sure it works with masking

lucidrains · May 2, 2023 · 220f104 · 220f104
1 parent 6a4106d
commit 220f104
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 4 deletions.
diff --git a/colt5_attention/triton_coor_descent.py b/colt5_attention/triton_coor_descent.py
@@ -232,12 +232,17 @@ def forward(
         x,
         n_iters,
         k,
-        eps
+        eps,
+        mask
     ):
         assert x.is_cuda
 
         batch, requires_grad = x.shape[0], x.requires_grad
 
+        if exists(mask):
+            mask_value = -torch.finfo(x.dtype).max
+            x = x.masked_fill(~mask, mask_value)
+
         x, shape = pack_one(x, '* n')
 
         n_rows, n_cols = x.shape
@@ -309,6 +314,14 @@ def backward(
             BLOCK_SIZE = BLOCK_SIZE
         )
 
-        return unpack_one(dx, shape, '* n'), None, None, None
+        return unpack_one(dx, shape, '* n'), None, None, None, None
 
-triton_coor_descent = _coor_descent.apply
+def triton_coor_descent(
+    s,
+    *,
+    n_iters,
+    k,
+    eps = 1e-1,
+    mask = None
+):
+    return _coor_descent.apply(s, n_iters, k, eps, mask)
diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'CoLT5-attention',
   packages = find_packages(),
-  version = '0.3.7',
+  version = '0.3.8',
   license='MIT',
   description = 'Conditionally Routed Attention',
   long_description_content_type = 'text/markdown',