Joeven — Learn Autonomous AI Agents

Python loads on first run

import math

X = [
    [1.0, 0.0, 0.0, 0.0],
    [0.0, 1.0, 0.0, 0.0],
    [0.8, 0.2, 0.0, 0.0],
]

def dot(a, b):
    return sum(x * y for x, y in zip(a, b))

def softmax(xs):
    m = max(xs)
    exps = [math.exp(x - m) for x in xs]
    s = sum(exps)
    return [e / s for e in exps]

def attention(X, causal=True):
    n = len(X)
    d = len(X[0])
    scale = math.sqrt(d)
    weights = []
    outs = []
    for i in range(n):
        scores = []
        for j in range(n):
            if causal and j > i:
                scores.append(-1e9)
            else:
                scores.append(dot(X[i], X[j]) / scale)
        w = softmax(scores)
        weights.append(w)
        mixed = [0.0] * d
        for j, wj in enumerate(w):
            for k in range(d):
                mixed[k] += wj * X[j][k]
        outs.append(mixed)
    return weights, outs

W, O = attention(X, causal=True)
print("causal weights (rows = query tokens):")
for i, row in enumerate(W):
    print("q" + str(i), [round(x, 3) for x in row])
print("out2", [round(x, 3) for x in O[2]])
Wb, _ = attention(X, causal=False)
print("bidirectional last row", [round(x, 3) for x in Wb[2]])

Run to execute this in your browser. Nothing is sent to a server.