(self, queries, keys, values, attn_mask)
| 145 | self.n_heads = n_heads |
| 146 | |
| 147 | def forward(self, queries, keys, values, attn_mask): |
| 148 | B, L, _ = queries.shape |
| 149 | _, S, _ = keys.shape |
| 150 | H = self.n_heads |
| 151 | |
| 152 | queries = self.query_projection(queries).view(B, L, H, -1) |
| 153 | keys = self.key_projection(keys).view(B, S, H, -1) |
| 154 | values = self.value_projection(values).view(B, S, H, -1) |
| 155 | |
| 156 | out, attn = self.inner_correlation( |
| 157 | queries, |
| 158 | keys, |
| 159 | values, |
| 160 | attn_mask |
| 161 | ) |
| 162 | out = out.view(B, L, -1) |
| 163 | |
| 164 | return self.out_projection(out), attn |
nothing calls this directly
no outgoing calls
no test coverage detected