Fix Ruff (#90)

stack-info: PR: #90, branch: drisspg/stack/3

Author: drisspg

attn_gym/masks/__init__.py

@@ -3,4 +3,4 @@
 from attn_gym.masks.prefix_lm import generate_prefix_lm_mask
 from attn_gym.masks.document_mask import generate_doc_mask_mod
 from attn_gym.masks.dilated_sliding_window import generate_dilated_sliding_window
-from attn_gym.masks.natten import generate_natten, generate_tiled_natten, generate_morton_natten
+from attn_gym.masks.natten import generate_natten, generate_tiled_natten, generate_morton_natten

attn_gym/masks/natten.py

@@ -42,6 +42,7 @@ def natten_mask_mod(
     natten_mask_mod.__name__ = f"natten_c{canvas_w}x{canvas_h}_k{kernel_w}x{kernel_h}"
     return natten_mask_mod
 
+
 def generate_tiled_natten(
     W: int,
     H: int,
@@ -68,7 +69,7 @@ def get_x_y_tiled(idx: IntTensor) -> Tuple[IntTensor, IntTensor]:
         t_x, t_y = t_id // (W // T_W), t_id % (W // T_W)
         t_offset = idx % (T_H * T_W)
         i_x, i_y = t_offset // T_W, t_offset % T_W
-        return t_x*T_W + i_x, t_y*T_H + i_y
+        return t_x * T_W + i_x, t_y * T_H + i_y
 
     def tiled_natten_mask(
         b: IntTensor,
@@ -87,6 +88,7 @@ def tiled_natten_mask(
     tiled_natten_mask.__name__ = f"tiled_natten_c{W}x{H}_k{K_W}x{K_H}_t{T_W}x{T_H}"
     return tiled_natten_mask
 
+
 def interleave_bits_32(x):
     """
     Interleave the bits of a 16-bit integer x, producing a 32-bit integer
@@ -99,6 +101,7 @@ def interleave_bits_32(x):
     x = (x | (x << 1)) & 0x55555555
     return x
 
+
 def morton_encode(x, y):
     """
     Encode 2D coordinates (x, y) into a Morton code (Z-order curve index).
@@ -112,6 +115,7 @@ def morton_encode(x, y):
     """
     return (interleave_bits_32(y) << 1) | interleave_bits_32(x)
 
+
 def deinterleave_bits_32(code):
     """
     Deinterleave bits to retrieve the original 16-bit integer.
@@ -123,6 +127,7 @@ def deinterleave_bits_32(code):
     code = (code | (code >> 8)) & 0x0000FFFF
     return code
 
+
 def morton_decode(code):
     """
     Decode a Morton code to retrieve the original 2D coordinates (x, y).
@@ -144,13 +149,14 @@ def generate_morton_natten(
     kernel_w: int,
     kernel_h: int,
 ) -> _mask_mod_signature:
-    """Generates a NATTEN attention mask with a given kernel size under morton curve layout. 
+    """Generates a NATTEN attention mask with a given kernel size under morton curve layout.
     Args:
         canvas_w: The width of the canvas.
         canvas_h: The height of the canvas.
         kernel_w: The width of the kernel.
         kernel_h: The height of the kernel.
     """
+
     def natten_mask_mod(
         b: IntTensor,
         h: IntTensor,
@@ -170,6 +176,7 @@ def natten_mask_mod(
     natten_mask_mod.__name__ = f"morton_natten_c{canvas_w}x{canvas_h}_k{kernel_w}x{kernel_h}"
     return natten_mask_mod
 
+
 def main(device: str = "cpu"):
     """Visualize the attention scores of NATTEN mask mod.
     Note: a more complete implementation of NATTEN would include support for kernel dilation.
@@ -204,8 +211,7 @@ def make_tensor():
         device=device,
         name=natten_mask.__name__,
     )
-    
-    
+
     tiled_natten_mask = generate_tiled_natten(
         W=CANVAS_WIDTH,
         H=CANVAS_HEIGHT,
@@ -222,8 +228,7 @@ def make_tensor():
         device=device,
         name=tiled_natten_mask.__name__,
     )
-    
-    
+
     morton_natten_mask = generate_morton_natten(
         canvas_w=CANVAS_WIDTH,
         canvas_h=CANVAS_HEIGHT,

test/test_natten.py

@@ -1,63 +1,69 @@
 import torch
-from torch.autograd import grad
 from torch.nn.attention.flex_attention import flex_attention, create_block_mask
 import pytest
-from functools import partial
 from attn_gym.masks import generate_natten, generate_tiled_natten, generate_morton_natten
 from attn_gym.masks.natten import morton_decode, morton_encode
 
 
-
 def run_natten(
-    mask = None,
-    encoder = None, 
-    decoder = None,
-    query = None, 
-    key = None,
-    value = None, 
-    gradOut = None,
+    mask=None,
+    encoder=None,
+    decoder=None,
+    query=None,
+    key=None,
+    value=None,
+    gradOut=None,
     print_mask=True,
 ):
     B, H, W, _, D = query.shape
     if decoder:
-        permuter_x, permuter_y = decoder(torch.arange(W*W))
-        permuter_index = permuter_x * W + permuter_y
-        q = query[:, :, permuter_x, permuter_y, :].clone().detach().requires_grad_(query.requires_grad)
+        permuter_x, permuter_y = decoder(torch.arange(W * W))
+        q = (
+            query[:, :, permuter_x, permuter_y, :]
+            .clone()
+            .detach()
+            .requires_grad_(query.requires_grad)
+        )
         k = key[:, :, permuter_x, permuter_y, :].clone().detach().requires_grad_(key.requires_grad)
-        v = value[:, :, permuter_x, permuter_y, :].clone().detach().requires_grad_(value.requires_grad)
+        v = (
+            value[:, :, permuter_x, permuter_y, :]
+            .clone()
+            .detach()
+            .requires_grad_(value.requires_grad)
+        )
         dO = gradOut[:, :, permuter_x, permuter_y, :]
-    else: 
+    else:
         q = query.flatten(2, 3).clone().detach().requires_grad_(query.requires_grad)
         k = key.flatten(2, 3).clone().detach().requires_grad_(key.requires_grad)
         v = value.flatten(2, 3).clone().detach().requires_grad_(value.requires_grad)
         dO = gradOut.flatten(2, 3)
-    block_mask = create_block_mask(mask, 1, 1, W*W, W*W, device=query.device)
+    block_mask = create_block_mask(mask, 1, 1, W * W, W * W, device=query.device)
     if print_mask:
         print(f"\nBlock Mask:\n{block_mask}")
-    
+
     flex_attention_compiled = torch.compile(flex_attention, dynamic=False)
     out = flex_attention_compiled(q, k, v, block_mask=block_mask)
-    
+
     out.backward(dO)
-    
-    if encoder: 
-        i_x = torch.arange(W)[:, None].broadcast_to(W, W).flatten() 
-        i_y = torch.arange(W)[None, :].broadcast_to(W, W).flatten() 
+
+    if encoder:
+        i_x = torch.arange(W)[:, None].broadcast_to(W, W).flatten()
+        i_y = torch.arange(W)[None, :].broadcast_to(W, W).flatten()
         depermuter = encoder(i_x, i_y)
         out = out[:, :, depermuter, :].reshape(B, H, W, W, D)
         q_grad = q.grad[:, :, depermuter, :].reshape(B, H, W, W, D)
         k_grad = k.grad[:, :, depermuter, :].reshape(B, H, W, W, D)
         v_grad = v.grad[:, :, depermuter, :].reshape(B, H, W, W, D)
         results = [out, q_grad, k_grad, v_grad]
     else:
-        out= out.reshape(B, H, W, W, D)
+        out = out.reshape(B, H, W, W, D)
         q_grad = q.grad.reshape(B, H, W, W, D)
         k_grad = k.grad.reshape(B, H, W, W, D)
         v_grad = v.grad.reshape(B, H, W, W, D)
         results = [out, q_grad, k_grad, v_grad]
-        
+
     del q, k, v, dO
-    
+
     return results
 
 
@@ -69,25 +75,21 @@ def test_natten_masks(
     K_W=13,
     T_W=8,
     print_mask=True,
-): 
-    query = torch.randn(
-        B, H, W, W, D, device="cuda", dtype=torch.float16, requires_grad=True
-    )
-    key = torch.randn(
-        B, H, W, W, D, device="cuda", dtype=torch.float16, requires_grad=True
-    )
-    value = torch.randn(
-        B, H, W, W, D, device="cuda", dtype=torch.float16, requires_grad=True
-    )
+):
+    query = torch.randn(B, H, W, W, D, device="cuda", dtype=torch.float16, requires_grad=True)
+    key = torch.randn(B, H, W, W, D, device="cuda", dtype=torch.float16, requires_grad=True)
+    value = torch.randn(B, H, W, W, D, device="cuda", dtype=torch.float16, requires_grad=True)
     gradOut = torch.randn(B, H, W, W, D, device="cuda", dtype=torch.float16)
-    
-    
+
     # Run naive NA
     naive_mask = generate_natten(W, W, K_W, K_W)
-    naive_results = run_natten(mask=naive_mask, query=query, key=key, value=value, gradOut=gradOut, print_mask=print_mask)
-    
+    naive_results = run_natten(
+        mask=naive_mask, query=query, key=key, value=value, gradOut=gradOut, print_mask=print_mask
+    )
+
     # Run tiled NA
     T_H = T_W
+
     def tiled_encoder(x, y):
         """
         Map 2-D coordinates to 1-D index for static tiles of T_H x T_W.
@@ -106,14 +108,33 @@ def tiled_decoder(idx):
         t_x, t_y = t_id // (W // T_W), t_id % (W // T_W)
         t_offset = idx % (T_H * T_W)
         i_x, i_y = t_offset // T_W, t_offset % T_W
-        return t_x*T_W + i_x, t_y*T_H + i_y
+        return t_x * T_W + i_x, t_y * T_H + i_y
+
     tiled_mask = generate_tiled_natten(W, W, K_W, K_W, T_W, T_H)
-    tiled_results = run_natten(mask=tiled_mask, encoder=tiled_encoder, decoder=tiled_decoder, query=query, key=key, value=value, gradOut=gradOut, print_mask=print_mask)
-    
+    tiled_results = run_natten(
+        mask=tiled_mask,
+        encoder=tiled_encoder,
+        decoder=tiled_decoder,
+        query=query,
+        key=key,
+        value=value,
+        gradOut=gradOut,
+        print_mask=print_mask,
+    )
+
     # Run morton NA
     morton_mask = generate_morton_natten(W, W, K_W, K_W)
-    morton_results = run_natten(mask=morton_mask, encoder=morton_encode, decoder=morton_decode, query=query, key=key, value=value, gradOut=gradOut, print_mask=print_mask)
-    
+    morton_results = run_natten(
+        mask=morton_mask,
+        encoder=morton_encode,
+        decoder=morton_decode,
+        query=query,
+        key=key,
+        value=value,
+        gradOut=gradOut,
+        print_mask=print_mask,
+    )
+
     for naive, tiled, morton in zip(naive_results, tiled_results, morton_results):
         torch.testing.assert_close(naive, tiled, atol=1e-1, rtol=1e-2)
         print("Tiled NATTEN: Correctness check passed ✅")
@@ -124,5 +145,6 @@ def tiled_decoder(idx):
     del query, key, value, gradOut, naive_results, tiled_results
     torch.cuda.empty_cache()
 
+
 if __name__ == "__main__":
-    pytest.main([__file__])
+    pytest.main([__file__])