Input shape: Time x Batch x Channel Args: key_padding_mask (ByteTensor, optional): mask to exclude keys that are pads, of shape `(batch, src_len)`, where padding elements are indicated by 1s. need_weights (bool, optional): return the a
(
self,
query, key, value,
key_padding_mask=None,
incremental_state=None,
need_weights=True,
static_kv=False,
attn_mask=None,
before_softmax=False,
need_head_weights=False,
enc_dec_attn_constraint_mask=None,
reset_attn_weight=None
)
| 234 | nn.init.xavier_normal_(self.bias_v) |
| 235 | |
| 236 | def forward( |
| 237 | self, |
| 238 | query, key, value, |
| 239 | key_padding_mask=None, |
| 240 | incremental_state=None, |
| 241 | need_weights=True, |
| 242 | static_kv=False, |
| 243 | attn_mask=None, |
| 244 | before_softmax=False, |
| 245 | need_head_weights=False, |
| 246 | enc_dec_attn_constraint_mask=None, |
| 247 | reset_attn_weight=None |
| 248 | ): |
| 249 | """Input shape: Time x Batch x Channel |
| 250 | |
| 251 | Args: |
| 252 | key_padding_mask (ByteTensor, optional): mask to exclude |
| 253 | keys that are pads, of shape `(batch, src_len)`, where |
| 254 | padding elements are indicated by 1s. |
| 255 | need_weights (bool, optional): return the attention weights, |
| 256 | averaged over heads (default: False). |
| 257 | attn_mask (ByteTensor, optional): typically used to |
| 258 | implement causal attention, where the mask prevents the |
| 259 | attention from looking forward in time (default: None). |
| 260 | before_softmax (bool, optional): return the raw attention |
| 261 | weights and values before the attention softmax. |
| 262 | need_head_weights (bool, optional): return the attention |
| 263 | weights for each head. Implies *need_weights*. Default: |
| 264 | return the average attention weights over all heads. |
| 265 | """ |
| 266 | if need_head_weights: |
| 267 | need_weights = True |
| 268 | |
| 269 | tgt_len, bsz, embed_dim = query.size() |
| 270 | assert embed_dim == self.embed_dim |
| 271 | assert list(query.size()) == [tgt_len, bsz, embed_dim] |
| 272 | |
| 273 | if self.enable_torch_version and incremental_state is None and not static_kv and reset_attn_weight is None: |
| 274 | if self.qkv_same_dim: |
| 275 | return F.multi_head_attention_forward(query, key, value, |
| 276 | self.embed_dim, self.num_heads, |
| 277 | self.in_proj_weight, |
| 278 | self.in_proj_bias, self.bias_k, self.bias_v, |
| 279 | self.add_zero_attn, self.dropout, |
| 280 | self.out_proj.weight, self.out_proj.bias, |
| 281 | self.training, key_padding_mask, need_weights, |
| 282 | attn_mask) |
| 283 | else: |
| 284 | return F.multi_head_attention_forward(query, key, value, |
| 285 | self.embed_dim, self.num_heads, |
| 286 | torch.empty([0]), |
| 287 | self.in_proj_bias, self.bias_k, self.bias_v, |
| 288 | self.add_zero_attn, self.dropout, |
| 289 | self.out_proj.weight, self.out_proj.bias, |
| 290 | self.training, key_padding_mask, need_weights, |
| 291 | attn_mask, use_separate_proj_weight=True, |
| 292 | q_proj_weight=self.q_proj_weight, |
| 293 | k_proj_weight=self.k_proj_weight, |
nothing calls this directly
no test coverage detected