(self, x)
| 140 | self.to_out = nn.Conv2d(hidden_dim, dim, 1) |
| 141 | |
| 142 | def forward(self, x): |
| 143 | b, c, h, w = x.shape |
| 144 | qkv = self.to_qkv(x) |
| 145 | q, k, v = rearrange(qkv, 'b (qkv heads c) h w -> qkv b heads c (h w)', heads=self.heads, qkv=3) |
| 146 | k = k.softmax(dim=-1) |
| 147 | context = torch.einsum('bhdn,bhen->bhde', k, v) |
| 148 | out = torch.einsum('bhde,bhdn->bhen', context, q) |
| 149 | out = rearrange(out, 'b heads c (h w) -> b (heads c) h w', heads=self.heads, h=h, w=w) |
| 150 | return self.to_out(out) |
| 151 | |
| 152 | |
| 153 | # gaussian diffusion trainer class |
nothing calls this directly
no outgoing calls
no test coverage detected