| 411 | |
| 412 | class NeighborhoodSelfAttentionBlock(nn.Module): |
| 413 | def __init__(self, d_model, d_head, cond_features, kernel_size, dropout=0.0): |
| 414 | super().__init__() |
| 415 | self.d_head = d_head |
| 416 | self.n_heads = d_model // d_head |
| 417 | self.kernel_size = kernel_size |
| 418 | self.norm = AdaRMSNorm(d_model, cond_features) |
| 419 | self.qkv_proj = apply_wd(Linear(d_model, d_model * 3, bias=False)) |
| 420 | self.scale = nn.Parameter(torch.full([self.n_heads], 10.0)) |
| 421 | self.pos_emb = AxialRoPE(d_head // 2, self.n_heads) |
| 422 | self.dropout = nn.Dropout(dropout) |
| 423 | self.out_proj = apply_wd(zero_init(Linear(d_model, d_model, bias=False))) |
| 424 | |
| 425 | def extra_repr(self): |
| 426 | return f"d_head={self.d_head}, kernel_size={self.kernel_size}" |