graduate_code/solutions/attention.py

"""G14: single-head scaled dot-product attention, not a trained transformer."""
import math


def softmax(scores, allowed=None):
    s = list(scores)
    mask = [True]*len(s) if allowed is None else list(allowed)
    if not s or len(mask) != len(s) or any(type(x) is not bool for x in mask) or not any(mask):
        raise ValueError("nonempty row with at least one allowed position required")
    if any(not math.isfinite(x) for x in s):
        raise ValueError("finite pre-mask scores required")
    top = max(v for v,ok in zip(s,mask) if ok)
    e = [math.exp(v-top) if ok else 0. for v,ok in zip(s,mask)]
    total = math.fsum(e)
    return [v/total for v in e]


def attention(Q,K,V,causal=False):
    n,m = len(Q),len(K)
    d = len(Q[0]) if n else 0
    dv = len(V[0]) if V else 0
    if (type(causal) is not bool or not 1 <= n <= 256 or not 1 <= m <= 256
            or not 1 <= d <= 128 or not 1 <= dv <= 128 or len(V) != m):
        raise ValueError("bounded nonempty matching matrices required")
    if causal and n != m:
        raise ValueError("this causal lab requires equal query/key lengths")
    for matrix,width in ((Q,d),(K,d),(V,dv)):
        if any(len(row) != width or any(not math.isfinite(x) for x in row) for row in matrix):
            raise ValueError("ragged or nonfinite matrix")
    out, weights = [],[]
    for i,q in enumerate(Q):
        scores = [math.fsum(a*b for a,b in zip(q,k))/math.sqrt(d) for k in K]
        p = softmax(scores,[j <= i for j in range(m)] if causal else None)
        weights.append(p)
        out.append([math.fsum(p[j]*V[j][v] for j in range(m)) for v in range(dv)])
    return out,weights


if __name__ == "__main__":
    print(attention([[0.],[0.]],[[1.],[2.]],[[2.],[8.]],causal=True))