A self-attention layer
In a model, the queries, keys and values are not given: each is a learned linear map of the same input. That is self-attention: the sequence attends to itself.
import torch
from torch import nn
d = 8
q = nn.Linear(d, d, bias=False)
# batch, positions, d
x = torch.randn(2, 5, d)
print(q(x).shape)
# torch.Size([2, 5, 8])The layer applies three such maps, scores every query against every key, masks the future, and mixes the values. Shapes stay (batch, positions, d), so layers can be stacked. With the causal mask, the output at a position depends only on that position and the ones before it, which is what lets a model learn to predict the next token from all of them at once.
Your task
Write the class SelfAttention(d), an nn.Module with three nn.Linear(d, d, bias=False) layers named q, k and v. Its forward(x) takes (batch, positions, d) and returns causal scaled dot-product attention of the three projections, of the same shape. You may compute it yourself or use F.scaled_dot_product_attention.
The explanation opens here once every check passes.