(
self,
hidden_size: int,
num_heads: int,
mlp_ratio: float = 4.0,
qk_scale: float | None = None,
backend='pytorch'
)
| 302 | """ |
| 303 | |
| 304 | def __init__( |
| 305 | self, |
| 306 | hidden_size: int, |
| 307 | num_heads: int, |
| 308 | mlp_ratio: float = 4.0, |
| 309 | qk_scale: float | None = None, |
| 310 | backend='pytorch' |
| 311 | ): |
| 312 | super().__init__() |
| 313 | self.hidden_dim = hidden_size |
| 314 | self.num_heads = num_heads |
| 315 | head_dim = hidden_size // num_heads |
| 316 | self.scale = qk_scale or head_dim**-0.5 |
| 317 | |
| 318 | self.mlp_hidden_dim = int(hidden_size * mlp_ratio) |
| 319 | # qkv and mlp_in |
| 320 | self.linear1 = nn.Linear(hidden_size, hidden_size * 3 + self.mlp_hidden_dim) |
| 321 | # proj and mlp_out |
| 322 | self.linear2 = nn.Linear(hidden_size + self.mlp_hidden_dim, hidden_size) |
| 323 | |
| 324 | self.norm = QKNorm(head_dim) |
| 325 | |
| 326 | self.hidden_size = hidden_size |
| 327 | self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) |
| 328 | |
| 329 | self.mlp_act = nn.GELU(approximate="tanh") |
| 330 | self.modulation = Modulation(hidden_size, double=False) |
| 331 | self.backend = backend |
| 332 | |
| 333 | def forward(self, x: Tensor, vec: Tensor, pe: Tensor, mask: Tensor = None) -> Tensor: |
| 334 | mod, _ = self.modulation(vec) |
nothing calls this directly
no test coverage detected