birder 0.3.3__py3-none-any.whl → 0.4.1__py3-none-any.whl

birder/net/_vit_configs.py

@@ -0,0 +1,466 @@
+"""
+ViT model configuration registrations
+
+This file contains *only* model variant definitions and their registration
+with the global model registry. The actual ViT implementation lives in vit.py.
+
+Naming:
+- All model names must follow the ViT / RoPE ViT naming convention documented in rope_vit_configs.py.
+"""
+
+from birder.model_registry import registry
+from birder.net.base import BaseNet
+
+TINY = {"num_layers": 12, "num_heads": 3, "hidden_dim": 192, "mlp_dim": 768, "drop_path_rate": 0.0}
+SMALL = {"num_layers": 12, "num_heads": 6, "hidden_dim": 384, "mlp_dim": 1536, "drop_path_rate": 0.0}
+MEDIUM = {"num_layers": 12, "num_heads": 8, "hidden_dim": 512, "mlp_dim": 2048, "drop_path_rate": 0.0}
+BASE = {"num_layers": 12, "num_heads": 12, "hidden_dim": 768, "mlp_dim": 3072, "drop_path_rate": 0.1}
+LARGE = {"num_layers": 24, "num_heads": 16, "hidden_dim": 1024, "mlp_dim": 4096, "drop_path_rate": 0.1}
+HUGE = {"num_layers": 32, "num_heads": 16, "hidden_dim": 1280, "mlp_dim": 5120, "drop_path_rate": 0.1}
+
+# From "Getting vit in Shape: Scaling Laws for Compute-Optimal Model Design"
+# Shape-optimized vision transformer (SoViT)
+SO150 = {
+    "num_layers": 18,
+    "num_heads": 16,
+    "hidden_dim": 896,  # Changed from 880 for RoPE divisibility
+    "mlp_dim": 2320,
+    "drop_path_rate": 0.1,
+}
+SO400 = {
+    "num_layers": 27,
+    "num_heads": 16,
+    "hidden_dim": 1152,
+    "mlp_dim": 4304,
+    "drop_path_rate": 0.1,
+}
+
+# From "Scaling Vision Transformers"
+GIANT = {"num_layers": 40, "num_heads": 16, "hidden_dim": 1408, "mlp_dim": 6144, "drop_path_rate": 0.1}
+GIGANTIC = {"num_layers": 48, "num_heads": 16, "hidden_dim": 1664, "mlp_dim": 8192, "drop_path_rate": 0.1}
+
+
+def register_vit_configs(vit: type[BaseNet]) -> None:
+    registry.register_model_config(
+        "vit_t32",
+        vit,
+        config={"patch_size": 32, **TINY},
+    )
+    registry.register_model_config(
+        "vit_t16",
+        vit,
+        config={"patch_size": 16, **TINY},
+    )
+    registry.register_model_config(
+        "vit_t14",
+        vit,
+        config={"patch_size": 14, **TINY},
+    )
+    registry.register_model_config(
+        "vit_s32",
+        vit,
+        config={"patch_size": 32, **SMALL},
+    )
+    registry.register_model_config(
+        "vit_s16",
+        vit,
+        config={"patch_size": 16, **SMALL},
+    )
+    registry.register_model_config(
+        "vit_s16_ls",
+        vit,
+        config={"patch_size": 16, **SMALL, "layer_scale_init_value": 1e-5},
+    )
+    registry.register_model_config(
+        "vit_s16_pn",
+        vit,
+        config={"patch_size": 16, **SMALL, "pre_norm": True, "norm_layer_eps": 1e-5},
+    )
+    registry.register_model_config(
+        "vit_s14",
+        vit,
+        config={"patch_size": 14, **SMALL},
+    )
+    registry.register_model_config(
+        "vit_m32",
+        vit,
+        config={"patch_size": 32, **MEDIUM},
+    )
+    registry.register_model_config(
+        "vit_m16",
+        vit,
+        config={"patch_size": 16, **MEDIUM},
+    )
+    registry.register_model_config(
+        "vit_m14",
+        vit,
+        config={"patch_size": 14, **MEDIUM},
+    )
+    registry.register_model_config(
+        "vit_b32",
+        vit,
+        config={"patch_size": 32, **BASE, "drop_path_rate": 0.0},  # Override the BASE definition
+    )
+    registry.register_model_config(
+        "vit_b16",
+        vit,
+        config={"patch_size": 16, **BASE},
+    )
+    registry.register_model_config(
+        "vit_b16_ls",
+        vit,
+        config={"patch_size": 16, **BASE, "layer_scale_init_value": 1e-5},
+    )
+    registry.register_model_config(
+        "vit_b16_qkn_ls",
+        vit,
+        config={"patch_size": 16, **BASE, "layer_scale_init_value": 1e-5, "qk_norm": True},
+    )
+    registry.register_model_config(
+        "vit_b16_pn_quick_gelu",
+        vit,
+        config={"patch_size": 16, **BASE, "pre_norm": True, "norm_layer_eps": 1e-5, "act_layer_type": "quick_gelu"},
+    )
+    registry.register_model_config(
+        "vit_b14",
+        vit,
+        config={"patch_size": 14, **BASE},
+    )
+    registry.register_model_config(
+        "vit_so150m_p14_avg",
+        vit,
+        config={"patch_size": 14, **SO150, "class_token": False},
+    )
+    registry.register_model_config(
+        "vit_so150m_p14_ap",
+        vit,
+        config={"patch_size": 14, **SO150, "class_token": False, "attn_pool_head": True},
+    )
+    registry.register_model_config(
+        "vit_l32",
+        vit,
+        config={"patch_size": 32, **LARGE},
+    )
+    registry.register_model_config(
+        "vit_l16",
+        vit,
+        config={"patch_size": 16, **LARGE},
+    )
+    registry.register_model_config(
+        "vit_l14",
+        vit,
+        config={"patch_size": 14, **LARGE},
+    )
+    registry.register_model_config(
+        "vit_l14_pn",
+        vit,
+        config={"patch_size": 14, **LARGE, "pre_norm": True, "norm_layer_eps": 1e-5},
+    )
+    registry.register_model_config(
+        "vit_l14_pn_quick_gelu",
+        vit,
+        config={"patch_size": 14, **LARGE, "pre_norm": True, "norm_layer_eps": 1e-5, "act_layer_type": "quick_gelu"},
+    )
+    registry.register_model_config(
+        "vit_so400m_p14_ap",
+        vit,
+        config={"patch_size": 14, **SO400, "class_token": False, "attn_pool_head": True},
+    )
+    registry.register_model_config(
+        "vit_h16",
+        vit,
+        config={"patch_size": 16, **HUGE},
+    )
+    registry.register_model_config(
+        "vit_h14",
+        vit,
+        config={"patch_size": 14, **HUGE},
+    )
+    registry.register_model_config(  # From "Scaling Vision Transformers"
+        "vit_g16",
+        vit,
+        config={"patch_size": 16, **GIANT},
+    )
+    registry.register_model_config(  # From "Scaling Vision Transformers"
+        "vit_g14",
+        vit,
+        config={"patch_size": 14, **GIANT},
+    )
+    registry.register_model_config(  # From "Scaling Vision Transformers"
+        "vit_gigantic14",
+        vit,
+        config={"patch_size": 14, **GIGANTIC},
+    )
+    registry.register_model_config(  # From "PaLI: A Jointly-Scaled Multilingual Language-Image Model"
+        "vit_e14",
+        vit,
+        config={
+            "patch_size": 14,
+            "num_layers": 56,
+            "num_heads": 16,
+            "hidden_dim": 1792,
+            "mlp_dim": 15360,
+            "drop_path_rate": 0.1,
+        },
+    )
+    registry.register_model_config(  # From "Scaling Language-Free Visual Representation Learning"
+        "vit_1b_p16",  # AKA vit_giant2 from DINOv2
+        vit,
+        config={
+            "patch_size": 16,
+            "num_layers": 40,
+            "num_heads": 24,
+            "hidden_dim": 1536,
+            "mlp_dim": 6144,
+            "drop_path_rate": 0.1,
+        },
+    )
+
+    # With registers
+    ####################
+
+    registry.register_model_config(
+        "vit_reg1_t32",
+        vit,
+        config={"patch_size": 32, **TINY, "num_reg_tokens": 1},
+    )
+    registry.register_model_config(
+        "vit_reg1_t16",
+        vit,
+        config={"patch_size": 16, **TINY, "num_reg_tokens": 1},
+    )
+    registry.register_model_config(
+        "vit_reg1_t14",
+        vit,
+        config={"patch_size": 14, **TINY, "num_reg_tokens": 1},
+    )
+    registry.register_model_config(
+        "vit_reg1_s32",
+        vit,
+        config={"patch_size": 32, **SMALL, "num_reg_tokens": 1},
+    )
+    registry.register_model_config(
+        "vit_reg1_s16",
+        vit,
+        config={"patch_size": 16, **SMALL, "num_reg_tokens": 1},
+    )
+    registry.register_model_config(
+        "vit_reg1_s16_ls",
+        vit,
+        config={"patch_size": 16, **SMALL, "layer_scale_init_value": 1e-5, "num_reg_tokens": 1},
+    )
+    registry.register_model_config(
+        "vit_reg1_s16_rms_ls",
+        vit,
+        config={
+            "patch_size": 16,
+            **SMALL,
+            "layer_scale_init_value": 1e-5,
+            "num_reg_tokens": 1,
+            "norm_layer_type": "RMSNorm",
+        },
+    )
+    registry.register_model_config(
+        "vit_reg1_s14",
+        vit,
+        config={"patch_size": 14, **SMALL, "num_reg_tokens": 1},
+    )
+    registry.register_model_config(
+        "vit_reg4_m32",
+        vit,
+        config={"patch_size": 32, **MEDIUM, "num_reg_tokens": 4},
+    )
+    registry.register_model_config(
+        "vit_reg4_m16",
+        vit,
+        config={"patch_size": 16, **MEDIUM, "num_reg_tokens": 4},
+    )
+    registry.register_model_config(
+        "vit_reg4_m16_rms_avg",
+        vit,
+        config={"patch_size": 16, **MEDIUM, "num_reg_tokens": 4, "class_token": False, "norm_layer_type": "RMSNorm"},
+    )
+    registry.register_model_config(
+        "vit_reg4_m14",
+        vit,
+        config={"patch_size": 14, **MEDIUM, "num_reg_tokens": 4},
+    )
+    registry.register_model_config(
+        "vit_reg4_b32",
+        vit,
+        config={"patch_size": 32, **BASE, "num_reg_tokens": 4, "drop_path_rate": 0.0},  # Override the BASE definition
+    )
+    registry.register_model_config(
+        "vit_reg4_b16",
+        vit,
+        config={"patch_size": 16, **BASE, "num_reg_tokens": 4},
+    )
+    registry.register_model_config(
+        "vit_reg4_b16_avg",
+        vit,
+        config={"patch_size": 16, **BASE, "num_reg_tokens": 4, "class_token": False},
+    )
+    registry.register_model_config(
+        "vit_reg4_b14",
+        vit,
+        config={"patch_size": 14, **BASE, "num_reg_tokens": 4},
+    )
+    registry.register_model_config(
+        "vit_reg8_b14_ap",
+        vit,
+        config={"patch_size": 14, **BASE, "num_reg_tokens": 8, "class_token": False, "attn_pool_head": True},
+    )
+    registry.register_model_config(
+        "vit_reg4_so150m_p16_avg",
+        vit,
+        config={"patch_size": 16, **SO150, "num_reg_tokens": 4, "class_token": False},
+    )
+    registry.register_model_config(
+        "vit_reg8_so150m_p16_swiglu_ap",
+        vit,
+        config={
+            "patch_size": 16,
+            **SO150,
+            "num_reg_tokens": 8,
+            "class_token": False,
+            "attn_pool_head": True,
+            "mlp_layer_type": "SwiGLU_FFN",
+        },
+    )
+    registry.register_model_config(
+        "vit_reg4_so150m_p14_avg",
+        vit,
+        config={"patch_size": 14, **SO150, "num_reg_tokens": 4, "class_token": False},
+    )
+    registry.register_model_config(
+        "vit_reg4_so150m_p14_ls",
+        vit,
+        config={"patch_size": 14, **SO150, "layer_scale_init_value": 1e-5, "num_reg_tokens": 4},
+    )
+    registry.register_model_config(
+        "vit_reg4_so150m_p14_ap",
+        vit,
+        config={"patch_size": 14, **SO150, "num_reg_tokens": 4, "class_token": False, "attn_pool_head": True},
+    )
+    registry.register_model_config(
+        "vit_reg4_so150m_p14_aps",
+        vit,
+        config={
+            "patch_size": 14,
+            **SO150,
+            "num_reg_tokens": 4,
+            "class_token": False,
+            "attn_pool_head": True,
+            "attn_pool_special_tokens": True,
+        },
+    )
+    registry.register_model_config(
+        "vit_reg8_so150m_p14_avg",
+        vit,
+        config={"patch_size": 14, **SO150, "num_reg_tokens": 8, "class_token": False},
+    )
+    registry.register_model_config(
+        "vit_reg8_so150m_p14_swiglu",
+        vit,
+        config={"patch_size": 14, **SO150, "num_reg_tokens": 8, "mlp_layer_type": "SwiGLU_FFN"},
+    )
+    registry.register_model_config(
+        "vit_reg8_so150m_p14_swiglu_avg",
+        vit,
+        config={"patch_size": 14, **SO150, "num_reg_tokens": 8, "class_token": False, "mlp_layer_type": "SwiGLU_FFN"},
+    )
+    registry.register_model_config(
+        "vit_reg8_so150m_p14_ap",
+        vit,
+        config={"patch_size": 14, **SO150, "num_reg_tokens": 8, "class_token": False, "attn_pool_head": True},
+    )
+    registry.register_model_config(
+        "vit_reg4_l32",
+        vit,
+        config={"patch_size": 32, **LARGE, "num_reg_tokens": 4},
+    )
+    registry.register_model_config(
+        "vit_reg4_l16",
+        vit,
+        config={"patch_size": 16, **LARGE, "num_reg_tokens": 4},
+    )
+    registry.register_model_config(
+        "vit_reg8_l16_avg",
+        vit,
+        config={"patch_size": 16, **LARGE, "num_reg_tokens": 8, "class_token": False},
+    )
+    registry.register_model_config(
+        "vit_reg8_l16_aps",
+        vit,
+        config={
+            "patch_size": 16,
+            **LARGE,
+            "num_reg_tokens": 8,
+            "class_token": False,
+            "attn_pool_head": True,
+            "attn_pool_special_tokens": True,
+        },
+    )
+    registry.register_model_config(
+        "vit_reg4_l14",
+        vit,
+        config={"patch_size": 14, **LARGE, "num_reg_tokens": 4},
+    )
+    registry.register_model_config(  # DeiT III style
+        "vit_reg4_l14_nps_ls",
+        vit,
+        config={
+            "pos_embed_special_tokens": False,
+            "patch_size": 14,
+            **LARGE,
+            "layer_scale_init_value": 1e-5,
+            "num_reg_tokens": 4,
+        },
+    )
+    registry.register_model_config(
+        "vit_reg8_l14_ap",
+        vit,
+        config={"patch_size": 14, **LARGE, "num_reg_tokens": 8, "class_token": False, "attn_pool_head": True},
+    )
+    registry.register_model_config(
+        "vit_reg8_l14_rms_ap",
+        vit,
+        config={
+            "patch_size": 14,
+            **LARGE,
+            "num_reg_tokens": 8,
+            "class_token": False,
+            "attn_pool_head": True,
+            "norm_layer_type": "RMSNorm",
+        },
+    )
+    registry.register_model_config(
+        "vit_reg8_so400m_p14_ap",
+        vit,
+        config={"patch_size": 14, **SO400, "num_reg_tokens": 8, "class_token": False, "attn_pool_head": True},
+    )
+    registry.register_model_config(
+        "vit_reg4_h16",
+        vit,
+        config={"patch_size": 16, **HUGE, "num_reg_tokens": 4},
+    )
+    registry.register_model_config(
+        "vit_reg4_h14",
+        vit,
+        config={"patch_size": 14, **HUGE, "num_reg_tokens": 4},
+    )
+    registry.register_model_config(  # From "Scaling Vision Transformers"
+        "vit_reg4_g16",
+        vit,
+        config={"patch_size": 16, **GIANT, "num_reg_tokens": 4},
+    )
+    registry.register_model_config(  # From "Scaling Vision Transformers"
+        "vit_reg4_g14",
+        vit,
+        config={"patch_size": 14, **GIANT, "num_reg_tokens": 4},
+    )
+    registry.register_model_config(  # From "Scaling Vision Transformers"
+        "vit_reg4_gigantic14",
+        vit,
+        config={"patch_size": 14, **GIGANTIC, "num_reg_tokens": 4},
+    )

birder/net/alexnet.py

@@ -27,17 +27,17 @@ class AlexNet(BaseNet):
         assert self.config is None, "config not supported"
 
         self.body = nn.Sequential(
-            nn.Conv2d(self.input_channels, 64, kernel_size=(11, 11), stride=(4, 4), padding=(2, 2), bias=True),
+            nn.Conv2d(self.input_channels, 64, kernel_size=(11, 11), stride=(4, 4), padding=(2, 2)),
             nn.ReLU(inplace=True),
             nn.MaxPool2d(kernel_size=(3, 3), stride=(2, 2), padding=(0, 0)),
-            nn.Conv2d(64, 192, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2), bias=True),
+            nn.Conv2d(64, 192, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2)),
             nn.ReLU(inplace=True),
             nn.MaxPool2d(kernel_size=(3, 3), stride=(2, 2), padding=(0, 0)),
-            nn.Conv2d(192, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=True),
+            nn.Conv2d(192, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
             nn.ReLU(inplace=True),
-            nn.Conv2d(384, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=True),
+            nn.Conv2d(384, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
             nn.ReLU(inplace=True),
-            nn.Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=True),
+            nn.Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
             nn.ReLU(inplace=True),
             nn.MaxPool2d(kernel_size=(3, 3), stride=(2, 2), padding=(0, 0)),
             nn.AdaptiveAvgPool2d(output_size=(6, 6)),

birder/net/base.py

@@ -5,6 +5,7 @@ from typing import Literal
 from typing import NotRequired
 from typing import Optional
 from typing import TypedDict
+from typing import overload
 
 import torch
 import torch.nn.functional as F
@@ -54,6 +55,30 @@ def make_divisible(v: float, divisor: int, min_value: Optional[int] = None) -> i
     return new_v
 
 
+@overload
+def normalize_out_indices(out_indices: None, num_layers: int) -> None: ...
+
+
+@overload
+def normalize_out_indices(out_indices: list[int], num_layers: int) -> list[int]: ...
+
+
+def normalize_out_indices(out_indices: Optional[list[int]], num_layers: int) -> Optional[list[int]]:
+    if out_indices is None:
+        return None
+
+    normalized_indices = []
+    for idx in out_indices:
+        if idx < 0:
+            idx = num_layers + idx
+        if idx < 0 or idx >= num_layers:
+            raise ValueError(f"out_indices contains invalid index for num_layers={num_layers}")
+
+        normalized_indices.append(idx)
+
+    return normalized_indices
+
+
 # class MiscNet(nn.Module):
 #     """
 #     Base class for general-purpose neural networks with automatic model registration
@@ -137,8 +162,8 @@ class BaseNet(nn.Module):
 
         self.dynamic_size = False
 
-        self.classifier: nn.Module
         self.embedding_size: int
+        self.classifier: nn.Module
 
     def create_classifier(self, embed_dim: Optional[int] = None) -> nn.Module:
         if self.num_classes == 0:
@@ -274,7 +299,7 @@ def pos_embedding_sin_cos_2d(
 ) -> torch.Tensor:
     # assert (dim % 4) == 0, "feature dimension must be multiple of 4 for sin-cos emb"
 
-    (y, x) = torch.meshgrid(torch.arange(h, device=device), torch.arange(w, device=device), indexing="ij")
+    y, x = torch.meshgrid(torch.arange(h, device=device), torch.arange(w, device=device), indexing="ij")
     omega = torch.arange(dim // 4, device=device) / (dim // 4 - 1)
     omega = 1.0 / (temperature**omega)
 
@@ -294,7 +319,7 @@ def interpolate_attention_bias(
     new_resolution: tuple[int, int],
     mode: Literal["bilinear", "bicubic"] = "bicubic",
 ) -> torch.Tensor:
-    (H, _) = attention_bias.size()
+    H, _ = attention_bias.size()
 
     # Interpolate
     orig_dtype = attention_bias.dtype

birder/net/biformer.py

@@ -8,6 +8,7 @@ Changes from original:
 * All attention types are in (B, C, H, W)
 * Using the newer Bi-Level Routing Attention implementation
 * Dynamic n_win size (image size // 32)
+* Stem bias term removed
 """
 
 # Reference license: Apache-2.0
@@ -29,7 +30,7 @@ from birder.net.base import DetectorBackbone
 
 
 def _grid2seq(x: torch.Tensor, region_size: tuple[int, int], num_heads: int) -> tuple[torch.Tensor, int, int]:
-    (B, C, H, W) = x.size()
+    B, C, H, W = x.size()
     region_h = H // region_size[0]
     region_w = W // region_size[1]
     x = x.view(B, num_heads, C // num_heads, region_h, region_size[0], region_w, region_size[1])
@@ -39,7 +40,7 @@ def _grid2seq(x: torch.Tensor, region_size: tuple[int, int], num_heads: int) ->
 
 
 def _seq2grid(x: torch.Tensor, region_h: int, region_w: int, region_size: tuple[int, int]) -> torch.Tensor:
-    (bs, n_head, _, _, head_dim) = x.size()
+    bs, n_head, _, _, head_dim = x.size()
     x = x.view(bs, n_head, region_h, region_w, region_size[0], region_size[1], head_dim)
     x = torch.einsum("bmhwpqd->bmdhpwq", x).reshape(
         bs, n_head * head_dim, region_h * region_size[0], region_w * region_size[1]
@@ -59,7 +60,7 @@ def regional_routing_attention_torch(
     auto_pad: bool,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     kv_region_size = region_size
-    (bs, n_head, q_nregion, topk) = region_graph.size()
+    bs, n_head, q_nregion, topk = region_graph.size()
 
     # Pad to deal with any input size
     q_pad_b = 0
@@ -67,13 +68,13 @@ def regional_routing_attention_torch(
     kv_pad_b = 0
     kv_pad_r = 0
     if auto_pad is True:
-        (_, _, h_q, w_q) = query.size()
+        _, _, h_q, w_q = query.size()
         q_pad_b = (region_size[0] - h_q % region_size[0]) % region_size[0]
         q_pad_r = (region_size[1] - w_q % region_size[1]) % region_size[1]
         if q_pad_b > 0 or q_pad_r > 0:
             query = F.pad(query, (0, q_pad_r, 0, q_pad_b))
 
-        (_, _, h_k, w_k) = key.size()
+        _, _, h_k, w_k = key.size()
         kv_pad_b = (kv_region_size[0] - h_k % kv_region_size[0]) % kv_region_size[0]
         kv_pad_r = (kv_region_size[1] - w_k % kv_region_size[1]) % kv_region_size[1]
         if kv_pad_r > 0 or kv_pad_b > 0:
@@ -86,12 +87,12 @@ def regional_routing_attention_torch(
         w_k = None
 
     # To sequence format
-    (query, q_region_h, q_region_w) = _grid2seq(query, region_size=region_size, num_heads=n_head)
-    (key, _, _) = _grid2seq(key, region_size=kv_region_size, num_heads=n_head)
-    (value, _, _) = _grid2seq(value, region_size=kv_region_size, num_heads=n_head)
+    query, q_region_h, q_region_w = _grid2seq(query, region_size=region_size, num_heads=n_head)
+    key, _, _ = _grid2seq(key, region_size=kv_region_size, num_heads=n_head)
+    value, _, _ = _grid2seq(value, region_size=kv_region_size, num_heads=n_head)
 
     # Gather key and values
-    (bs, n_head, kv_nregion, kv_region_size, head_dim) = key.size()
+    bs, n_head, kv_nregion, kv_region_size, head_dim = key.size()
     broadcasted_region_graph = region_graph.view(bs, n_head, q_nregion, topk, 1, 1).expand(
         -1, -1, -1, -1, kv_region_size, head_dim
     )
@@ -145,12 +146,12 @@ class BiLevelRoutingAttention(nn.Module):
         self.output_linear = nn.Conv2d(dim, dim, kernel_size=(1, 1), stride=(1, 1), padding=(0, 0))
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        (_, _, H, W) = x.size()
+        _, _, H, W = x.size()
         region_size = (H // self.n_win_h, W // self.n_win_w)
 
         # Linear projection
         qkv = self.qkv_linear(x)
-        (q, k, v) = qkv.chunk(3, dim=1)
+        q, k, v = qkv.chunk(3, dim=1)
 
         # Region-to-region routing
         q_r = F.avg_pool2d(  # pylint: disable=not-callable
@@ -162,11 +163,11 @@ class BiLevelRoutingAttention(nn.Module):
         q_r = q_r.permute(0, 2, 3, 1).flatten(1, 2)  # (n, (hw), c)
         k_r = k_r.flatten(2, 3)  # (n, c, (hw))
         a_r = q_r @ k_r
-        (_, idx_r) = torch.topk(a_r, k=self.topk, dim=-1)
+        _, idx_r = torch.topk(a_r, k=self.topk, dim=-1)
         idx_r = idx_r.unsqueeze_(1).expand(-1, self.num_heads, -1, -1)
 
         # Token to token attention
-        (output, _) = regional_routing_attention_torch(
+        output, _ = regional_routing_attention_torch(
             q, k, v, scale=self.scale, region_graph=idx_r, region_size=region_size, auto_pad=True
         )
 
@@ -189,12 +190,12 @@ class Attention(nn.Module):
         self.proj_drop = nn.Dropout(proj_drop)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        (B, C, H, W) = x.size()
+        B, C, H, W = x.size()
         x = x.permute(0, 2, 3, 1).reshape(B, H * W, C)
 
         N = H * W
         qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
-        (q, k, v) = qkv.unbind(0)
+        q, k, v = qkv.unbind(0)
 
         x = F.scaled_dot_product_attention(  # pylint: disable=not-callable
             q, k, v, dropout_p=self.attn_drop.p if self.training else 0.0, scale=self.scale
@@ -236,8 +237,8 @@ class AttentionLePE(nn.Module):
         )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        (B, C, H, W) = x.size()
-        (q, k, v) = self.qkv(x).chunk(3, dim=1)
+        B, C, H, W = x.size()
+        q, k, v = self.qkv(x).chunk(3, dim=1)
 
         attn = q.view(B, self.num_heads, self.head_dim, H * W).transpose(-1, -2) @ k.view(
             B, self.num_heads, self.head_dim, H * W