AppFolio Property Manager may use attention-based models to represent sequences such as maintenance updates or resident communications. Implement a single self-attention layer from scratch using only Python's standard library.
Given token embeddings and projection matrices, compute queries, keys, and values, then return scaled dot-product attention outputs. Do not use NumPy, PyTorch, TensorFlow, or other numerical libraries.
Implement self_attention(tokens, w_q, w_k, w_v, causal), where tokens is an n x d_model matrix and each projection matrix is d_model x d_head. Compute Q = tokens @ w_q, K = tokens @ w_k, and V = tokens @ w_v. For every token pair (i, j), compute:
score(i, j) = dot(Q[i], K[j]) / sqrt(d_head)
Apply a row-wise softmax to the scores. When causal is true, position i may attend only to positions j <= i. Return the resulting n x d_head matrix, where each output row is the weighted sum of value rows.
The softmax must be numerically stable by subtracting the row maximum before exponentiation.
def self_attention(tokens, w_q, w_k, w_v, causal):