(query, key, value)
| 18 | return out |
| 19 | |
| 20 | def attention(query, key, value): |
| 21 | scale = 1 / query.shape[-1] ** 0.5 |
| 22 | query = query * scale |
| 23 | attn = query @ key.transpose(-2, -1) |
| 24 | attn = attn.softmax(-1) |
| 25 | attn = torch.nn.functional.dropout(attn, 0) |
| 26 | return attn @ value |
| 27 | |
| 28 | def test_standard_attention(seq_len, hidden_size, num_heads, batch_size=8, qkv=None): |
| 29 | # set random seed |
no outgoing calls
no test coverage detected