(self, query, values, attn_mask=None)
| 23 | self.softmax = torch.nn.Softmax(dim=-1) |
| 24 | |
| 25 | def forward(self, query, values, attn_mask=None): |
| 26 | # query shape: batch x query_size |
| 27 | # values shape: batch x num values x value_size |
| 28 | |
| 29 | # attn_logits shape: batch x num values |
| 30 | attn_logits = self.pointer(query, values, attn_mask) |
| 31 | # attn_logits shape: batch x num values |
| 32 | attn = self.softmax(attn_logits) |
| 33 | # output shape: batch x 1 x value_size |
| 34 | output = torch.bmm(attn.unsqueeze(1), values) |
| 35 | output = output.squeeze(1) |
| 36 | return output, attn |
| 37 | |
| 38 | |
| 39 | @registry.register('pointer', 'sdp') |
nothing calls this directly
no outgoing calls
no test coverage detected