(
self,
channels: int,
num_heads: int,
mlp_ratio: float = 4.0,
attn_mode: Literal["full", "windowed"] = "full",
window_size: Optional[int] = None,
shift_window: Optional[int] = None,
use_checkpoint: bool = False,
use_rope: bool = False,
qk_rms_norm: bool = False,
qkv_bias: bool = True,
ln_affine: bool = False,
)
| 64 | Transformer block (MSA + FFN). |
| 65 | """ |
| 66 | def __init__( |
| 67 | self, |
| 68 | channels: int, |
| 69 | num_heads: int, |
| 70 | mlp_ratio: float = 4.0, |
| 71 | attn_mode: Literal["full", "windowed"] = "full", |
| 72 | window_size: Optional[int] = None, |
| 73 | shift_window: Optional[int] = None, |
| 74 | use_checkpoint: bool = False, |
| 75 | use_rope: bool = False, |
| 76 | qk_rms_norm: bool = False, |
| 77 | qkv_bias: bool = True, |
| 78 | ln_affine: bool = False, |
| 79 | ): |
| 80 | super().__init__() |
| 81 | self.use_checkpoint = use_checkpoint |
| 82 | self.norm1 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6) |
| 83 | self.norm2 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6) |
| 84 | self.attn = MultiHeadAttention( |
| 85 | channels, |
| 86 | num_heads=num_heads, |
| 87 | attn_mode=attn_mode, |
| 88 | window_size=window_size, |
| 89 | shift_window=shift_window, |
| 90 | qkv_bias=qkv_bias, |
| 91 | use_rope=use_rope, |
| 92 | qk_rms_norm=qk_rms_norm, |
| 93 | ) |
| 94 | self.mlp = FeedForwardNet( |
| 95 | channels, |
| 96 | mlp_ratio=mlp_ratio, |
| 97 | ) |
| 98 | |
| 99 | def _forward(self, x: torch.Tensor) -> torch.Tensor: |
| 100 | h = self.norm1(x) |
nothing calls this directly
no test coverage detected