4 hidden checks

Implement the attention function at the heart of every transformer.

  • softmax(x, axis=-1): numerically stable softmax along an axis. It must handle -np.inf entries (they get weight 0).
  • attention(Q, K, V, causal=False): return (out, weights) where weights = softmax(Q @ K.T / sqrt(d_k)) and out = weights @ V. With causal=True, set scores above the diagonal to -np.inf first (use np.triu(..., k=1)). Assume Q and K have the same number of rows when causal.
checks
  • softmax is stable and handles -infnot run
  • attention matches the formulanot run
  • causal attention never looks aheadnot run
  • changing a future token doesn't change earlier outputsnot run
python sleeps until you press run

Press Run & check. Your code runs first, then the hidden checks on the left.