| 38 | |
| 39 | @registry.register('pointer', 'sdp') |
| 40 | class ScaledDotProductPointer(torch.nn.Module): |
| 41 | def __init__(self, query_size, key_size): |
| 42 | super().__init__() |
| 43 | self.query_proj = torch.nn.Linear(query_size, key_size) |
| 44 | self.temp = np.power(key_size, 0.5) |
| 45 | |
| 46 | def forward(self, query, keys, attn_mask=None): |
| 47 | # query shape: batch x query_size |
| 48 | # keys shape: batch x num keys x key_size |
| 49 | |
| 50 | # proj_query shape: batch x key_size x 1 |
| 51 | proj_query = self.query_proj(query).unsqueeze(2) |
| 52 | |
| 53 | # attn_logits shape: batch x num keys |
| 54 | attn_logits = torch.bmm(keys, proj_query).squeeze(2) / self.temp |
| 55 | maybe_mask(attn_logits, attn_mask) |
| 56 | return attn_logits |
| 57 | |
| 58 | |
| 59 | @registry.register('attention', 'sdp') |