Softmax: stable probabilities and gradients

import numpy as np

def softmax(x):
    x = np.asarray(x, dtype=float)
    if x.size == 0 or not np.isfinite(x).all():
        raise ValueError('nonempty finite logits required')
    z = np.exp(x - x.max(axis=-1, keepdims=True))
    return z / z.sum(axis=-1, keepdims=True)

def attention(q, k, v, causal=False):
    if q.shape[-1] != k.shape[-1] or k.shape[-2] != v.shape[-2]:
        raise ValueError('incompatible shapes')
    scores = q @ k.swapaxes(-1, -2) / np.sqrt(q.shape[-1])
    if causal:
        if q.shape[-2] != k.shape[-2]:
            raise ValueError('this teaching mask requires equal sequence lengths')
        mask = np.tril(np.ones(scores.shape[-2:], dtype=bool))
        scores = np.where(mask, scores, -np.inf)
    z = np.exp(scores - scores.max(axis=-1, keepdims=True))
    weights = z / z.sum(axis=-1, keepdims=True)
    return weights @ v

def multi_head(x, wq, wk, wv, wo, heads, causal=True):
    b, t, d = x.shape
    if heads < 1 or d % heads: raise ValueError('heads must divide model dimension')
    def split(w):
        return (x @ w).reshape(b, t, heads, d // heads).transpose(0, 2, 1, 3)
    y = attention(split(wq), split(wk), split(wv), causal)
    return y.transpose(0, 2, 1, 3).reshape(b, t, d) @ wo

def kmeans(x, centers, max_iter=100, tol=1e-8):
    x, c = np.asarray(x, float), np.array(centers, float, copy=True)
    if x.ndim!=2 or c.ndim!=2 or not len(x) or not len(c) or x.shape[1]!=c.shape[1] or max_iter<1 or not np.isfinite(x).all() or not np.isfinite(c).all():
        raise ValueError('invalid data or centers')
    for _ in range(max_iter):
        labels = ((x[:,None,:]-c[None,:,:])**2).sum(-1).argmin(-1)
        new = np.array([x[labels==j].mean(0) if np.any(labels==j) else c[j] for j in range(len(c))])
        shift = np.max(np.linalg.norm(new-c,axis=1))
        c = new
        if shift <= tol: break
    labels = ((x[:,None,:]-c[None,:,:])**2).sum(-1).argmin(-1)
    return c, labels, ((x-c[labels])**2).sum()

class Adam:
    def __init__(self, shape, lr=.001, b1=.9, b2=.999, eps=1e-8):
        self.m=np.zeros(shape); self.v=np.zeros(shape); self.t=0
        self.lr,self.b1,self.b2,self.eps=lr,b1,b2,eps
    def step(self, p, g):
        if p.shape!=self.m.shape or g.shape!=p.shape or not np.isfinite(g).all(): raise ValueError('invalid gradient')
        self.t+=1
        self.m=self.b1*self.m+(1-self.b1)*g
        self.v=self.b2*self.v+(1-self.b2)*g*g
        mh=self.m/(1-self.b1**self.t); vh=self.v/(1-self.b2**self.t)
        return p-self.lr*mh/(np.sqrt(vh)+self.eps)

def beam_search(next_log_probs, width=2, max_steps=4, eos='EOS'):
    if width<1 or max_steps<1: raise ValueError('positive limits required')
    beams=[((),0.,False)]
    for _ in range(max_steps):
        candidates=[]
        for seq,score,done in beams:
            if done: candidates.append((seq,score,True)); continue
            for token,lp in next_log_probs(seq).items():
                if not np.isfinite(lp) or lp>0: raise ValueError('finite log probabilities required')
                candidates.append((seq+(token,),score+lp,token==eos))
        if not candidates: raise ValueError('no candidates')
        beams=sorted(candidates,key=lambda x:(-x[1],x[0]))[:width]
        if all(x[2] for x in beams): break
    return beams

if __name__=='__main__':
    assert np.allclose(softmax([1000,1000]),[.5,.5])
    assert np.allclose(softmax([1,2]),softmax([101,102]))
    x=np.arange(24,dtype=float).reshape(1,3,8)/24; eye=np.eye(8)
    y=multi_head(x,eye,eye,eye,eye,2); changed=x.copy(); changed[:,2]+=10
    assert np.allclose(y[:,:2],multi_head(changed,eye,eye,eye,eye,2)[:,:2])
    assert np.allclose(multi_head(x,eye,eye,eye,eye,1),attention(x,x,x,True))
    c,l,s=kmeans([[0],[2],[8],[10]],[[0],[10]])
    assert np.allclose(c,[[1],[9]]) and s==4
    c,l,s=kmeans([[0],[0]],[[0],[8]]); assert np.isfinite(c).all()
    opt=Adam((1,)); assert np.allclose(opt.step(np.array([1.]),np.array([2.])),[.999])
    assert opt.t==1
    def probs(seq):
        p={'A':.6,'B':.4} if not seq else ({'EOS':.5,'X':.5} if seq[-1]=='A' else {'EOS':.9,'X':.1})
        return {k:np.log(v) for k,v in p.items()}
    assert beam_search(probs,max_steps=2)[0][0]==('B','EOS')
    assert beam_search(probs,width=1,max_steps=2)[0][0]==('A','EOS')
    print('PASS: 10 algorithm invariants (synthetic fixtures)')

Hint: Cancel the common factor.