Attention is the core of transformers, the models behind today's language AI. Each position asks a question (a query), every position offers a key and a value, and each position gets back a mix of the values, weighted by how well its query matches each key.

import math
import torch

q, k, v = torch.randn(3, 4, 8)
d = q.shape[-1]
scores = q @ k.T / math.sqrt(d)
# positions x positions, rows sum to 1
weights = torch.softmax(scores, dim=-1)
output = weights @ v
print(weights.sum(dim=-1))
print(output.shape)

Dividing by sqrt(d), where d is the query size, keeps the scores from growing with the size and the softmax from becoming too sharp. A causal mask stops a position from looking at later ones: their scores are set to minus infinity before the softmax, so their weights become 0. That is how a model that predicts the next word is kept from seeing it.

PyTorch has this built in as F.scaled_dot_product_attention, with import torch.nn.functional as F.

Write attention(q, k, v, causal=False) for tensors of shape (positions, d): compute scaled dot-product attention yourself, with torch.softmax, applying a causal mask when asked. Return the pair (output, weights).

The explanation opens here once every check passes.