| 159 | ) |
| 160 | |
| 161 | def forward(self, x, context=None, mask=None): |
| 162 | h = self.heads |
| 163 | |
| 164 | q = self.to_q(x) |
| 165 | context = default(context, x) |
| 166 | k = self.to_k(context) |
| 167 | v = self.to_v(context) |
| 168 | |
| 169 | q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v)) |
| 170 | |
| 171 | # force cast to fp32 to avoid overflowing |
| 172 | if _ATTN_PRECISION =="fp32": |
| 173 | with torch.autocast(enabled=False, device_type = 'cuda'): |
| 174 | q, k = q.float(), k.float() |
| 175 | sim = einsum('b i d, b j d -> b i j', q, k) * self.scale |
| 176 | else: |
| 177 | sim = einsum('b i d, b j d -> b i j', q, k) * self.scale |
| 178 | |
| 179 | del q, k |
| 180 | |
| 181 | if exists(mask): |
| 182 | mask = rearrange(mask, 'b ... -> b (...)') |
| 183 | max_neg_value = -torch.finfo(sim.dtype).max |
| 184 | mask = repeat(mask, 'b j -> (b h) () j', h=h) |
| 185 | sim.masked_fill_(~mask, max_neg_value) |
| 186 | |
| 187 | # attention, what we cannot get enough of |
| 188 | sim = sim.softmax(dim=-1) |
| 189 | |
| 190 | out = einsum('b i j, b j d -> b i d', sim, v) |
| 191 | out = rearrange(out, '(b h) n d -> b n (h d)', h=h) |
| 192 | return self.to_out(out) |
| 193 | |
| 194 | |
| 195 | class MemoryEfficientCrossAttention(nn.Module): |