Add option for cauasl (#107)

Author: drisspg

attn_gym/masks/document_mask.py

@@ -4,7 +4,7 @@
 
 import torch
 from torch import Tensor
-from torch.nn.attention.flex_attention import _mask_mod_signature
+from torch.nn.attention.flex_attention import _mask_mod_signature, noop_mask
 from attn_gym.masks import causal_mask
 
 
@@ -59,7 +59,7 @@ def doc_mask_mod(b, h, q_idx, kv_idx):
     return doc_mask_mod
 
 
-def main(device: str = "cpu"):
+def main(device: str = "cpu", causal: bool = True):
     """Visualize the attention scores of document causal mask mod.
 
     Args:
@@ -93,7 +93,11 @@ def make_tensor():
         return torch.ones(B, H, SEQ_LEN, HEAD_DIM, device=device)
 
     query, key = make_tensor(), make_tensor()
-    document_causal_mask = generate_doc_mask_mod(causal_mask, offsets)
+    if causal:
+        base_mask_mod = causal_mask
+    else:
+        base_mask_mod = noop_mask
+    document_causal_mask = generate_doc_mask_mod(base_mask_mod, offsets)
 
     visualize_attention_scores(
         query,