(
self,
dim: int,
num_heads: int,
mlp_ratio: float = 4.0,
qkv_bias: bool = False,
proj_bias: bool = True,
ffn_bias: bool = True,
drop: float = 0.0,
attn_drop: float = 0.0,
init_values=None,
drop_path: float = 0.0,
act_layer: Callable[..., nn.Module] = nn.GELU,
norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
attn_class: Callable[..., nn.Module] = Attention,
ffn_layer: Callable[..., nn.Module] = Mlp,
qk_norm: bool=False,
rope=None
)
| 258 | |
| 259 | class BlockRope(nn.Module): |
| 260 | def __init__( |
| 261 | self, |
| 262 | dim: int, |
| 263 | num_heads: int, |
| 264 | mlp_ratio: float = 4.0, |
| 265 | qkv_bias: bool = False, |
| 266 | proj_bias: bool = True, |
| 267 | ffn_bias: bool = True, |
| 268 | drop: float = 0.0, |
| 269 | attn_drop: float = 0.0, |
| 270 | init_values=None, |
| 271 | drop_path: float = 0.0, |
| 272 | act_layer: Callable[..., nn.Module] = nn.GELU, |
| 273 | norm_layer: Callable[..., nn.Module] = nn.LayerNorm, |
| 274 | attn_class: Callable[..., nn.Module] = Attention, |
| 275 | ffn_layer: Callable[..., nn.Module] = Mlp, |
| 276 | qk_norm: bool=False, |
| 277 | rope=None |
| 278 | ) -> None: |
| 279 | super().__init__() |
| 280 | # print(f"biases: qkv: {qkv_bias}, proj: {proj_bias}, ffn: {ffn_bias}") |
| 281 | self.norm1 = norm_layer(dim) |
| 282 | self.attn = attn_class( |
| 283 | dim, |
| 284 | num_heads=num_heads, |
| 285 | qkv_bias=qkv_bias, |
| 286 | proj_bias=proj_bias, |
| 287 | attn_drop=attn_drop, |
| 288 | proj_drop=drop, |
| 289 | qk_norm=qk_norm, |
| 290 | rope=rope |
| 291 | ) |
| 292 | |
| 293 | self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity() |
| 294 | self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() |
| 295 | |
| 296 | self.norm2 = norm_layer(dim) |
| 297 | mlp_hidden_dim = int(dim * mlp_ratio) |
| 298 | self.mlp = ffn_layer( |
| 299 | in_features=dim, |
| 300 | hidden_features=mlp_hidden_dim, |
| 301 | act_layer=act_layer, |
| 302 | drop=drop, |
| 303 | bias=ffn_bias, |
| 304 | ) |
| 305 | self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity() |
| 306 | self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() |
| 307 | |
| 308 | self.sample_drop_ratio = drop_path |
| 309 | |
| 310 | def forward(self, x: Tensor, xpos=None) -> Tensor: |
| 311 | def attn_residual_func(x: Tensor) -> Tensor: |
nothing calls this directly
no test coverage detected