(self, x, attn_mask=None)
| 137 | nn.init.constant_(m.weight, 1.0) |
| 138 | |
| 139 | def forward(self, x, attn_mask=None): |
| 140 | |
| 141 | pos_embed = get_abs_pos(self.pos_embed, x.size(1)) |
| 142 | |
| 143 | x = self.kv_proj(x) |
| 144 | x = self.ln_kv(x).permute(1, 0, 2) |
| 145 | |
| 146 | N = x.shape[1] |
| 147 | q = self.ln_q(self.query) |
| 148 | out = self.attn( |
| 149 | self._repeat(q, N) + self.pos_embed.unsqueeze(1), |
| 150 | x + pos_embed.unsqueeze(1), |
| 151 | x, |
| 152 | attn_mask=attn_mask)[0] |
| 153 | return out.permute(1, 0, 2) |
| 154 | |
| 155 | def _repeat(self, query, N: int): |
| 156 | return query.unsqueeze(1).repeat(1, N, 1) |
nothing calls this directly
no test coverage detected