1import json, math, random, time
  2import numpy as np
  3
  4
  5def empty(k):
  6    return np.zeros(k, dtype=np.float64), np.zeros((k, k), dtype=np.float64), np.zeros(k, dtype=np.float64)
  7
  8
  9def update(state, v):
 10    u, A, Q = state
 11    v = np.asarray(v, dtype=np.float64)
 12    return (u + v,
 13            A + 0.5 * (np.outer(u, v) - np.outer(v, u)),
 14            Q + v * v)
 15
 16
 17def compose(left, right):
 18    u, A, Q = left
 19    v, B, R = right
 20    return (u + v,
 21            A + B + 0.5 * (np.outer(u, v) - np.outer(v, u)),
 22            Q + R)
 23
 24
 25def sequence(xs):
 26    s = empty(len(xs[0]))
 27    for x in xs:
 28        s = update(s, x)
 29    return s
 30
 31
 32def maxdiff(a, b):
 33    return max(float(np.max(np.abs(x - y))) for x, y in zip(a, b))
 34
 35
 36def math_checks():
 37    # Prediction 1: for orthogonal increments a e1 then b e2, A12=ab/2;
 38    # reversing order gives -ab/2 and therefore cancels exactly.
 39    reversal = []
 40    for a in [0.25, 0.5, 1.0, 2.0, 4.0]:
 41        forward = sequence([np.array([a, 0.0]), np.array([0.0, 1.0])])[1][0, 1]
 42        reverse = sequence([np.array([0.0, 1.0]), np.array([a, 0.0])])[1][0, 1]
 43        reversal.append({"a": a, "predicted_abs": a / 2.0,
 44                         "observed_abs": abs(forward),
 45                         "reverse_sum": forward + reverse})
 46
 47    # Prediction 2: n repeated ordered pairs accumulate linearly: A12=n/2.
 48    accumulation = []
 49    for n in range(1, 9):
 50        xs = []
 51        for _ in range(n):
 52            xs.extend([np.array([1.0, 0.0]), np.array([0.0, 1.0])])
 53        observed = sequence(xs)[1][0, 1]
 54        accumulation.append({"n": n, "predicted": n / 2.0, "observed": observed})
 55
 56    # Prediction 3: under x -> lambda*x, u has degree one while A,Q have degree two.
 57    rng = np.random.default_rng(4)
 58    xs = rng.normal(size=(7, 3))
 59    base = sequence(xs)
 60    dilation = []
 61    for lam in [0.25, 0.5, 1.0, 2.0, 4.0]:
 62        got = sequence(lam * xs)
 63        dilation.append({
 64            "lambda": lam,
 65            "u_ratio_to_lambda": float(np.linalg.norm(got[0]) / (np.linalg.norm(base[0]) * lam)),
 66            "A_ratio_to_lambda2": float(np.linalg.norm(got[1]) / (np.linalg.norm(base[1]) * lam * lam)),
 67            "Q_ratio_to_lambda2": float(np.linalg.norm(got[2]) / (np.linalg.norm(base[2]) * lam * lam)),
 68        })
 69
 70    # Associativity/chunking prediction: recursively summarized chunks equal token scan.
 71    rng = np.random.default_rng(11)
 72    xs = rng.normal(size=(13, 5))
 73    token = sequence(xs)
 74    chunks = [sequence(xs[:4]), sequence(xs[4:9]), sequence(xs[9:])]
 75    merged = compose(compose(chunks[0], chunks[1]), chunks[2])
 76    return {"reversal_sweep": reversal, "linear_accumulation": accumulation,
 77            "dilation_sweep": dilation, "chunk_max_abs_error": maxdiff(token, merged)}
 78
 79
 80def benchmark_impl(seed=123, epochs=80):
 81    # Task deliberately requires order: label is whether the first symbol precedes
 82    # the second symbol, while the multiset of symbols is identical.
 83    import torch
 84    import torch.nn as nn
 85    torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)
 86    device = "cuda" if torch.cuda.is_available() else "cpu"
 87    try:
 88        torch.zeros(1, device=device)
 89    except Exception:
 90        device = "cpu"
 91    ntrain, ntest, length, vocab, k = 768, 256, 12, 8, 8
 92    rng = np.random.default_rng(seed)
 93    def make(n):
 94        x = np.zeros((n, length), dtype=np.int64); y=np.zeros(n, dtype=np.int64)
 95        for i in range(n):
 96            a,b=rng.choice(vocab, 2, replace=False)
 97            x[i,:]=rng.integers(0,vocab,length)
 98            pos=rng.choice(length,2,replace=False); x[i,pos[0]]=a; x[i,pos[1]]=b
 99            y[i]=int(pos[0] < pos[1])
100        return torch.tensor(x), torch.tensor(y)
101    tr_x,tr_y=make(ntrain); te_x,te_y=make(ntest)
102
103    class Additive(nn.Module):
104        def __init__(self):
105            super().__init__(); self.emb=nn.Embedding(vocab,k); self.head=nn.Linear(k,2)
106        def forward(self,x): return self.head(self.emb(x).sum(1))
107
108    class Levy(nn.Module):
109        def __init__(self):
110            super().__init__(); self.emb=nn.Embedding(vocab,k); self.head=nn.Linear(k+k*(k-1)//2+k,2)
111        def forward(self,x):
112            z=torch.tanh(self.emb(x)); u=torch.zeros(x.size(0),k,device=x.device); A=torch.zeros(x.size(0),k,k,device=x.device); q=torch.zeros_like(u)
113            for t in range(length):
114                v=z[:,t]; A=A+0.5*(u[:,:,None]*v[:,None,:]-v[:,:,None]*u[:,None,:]); u=u+v; q=q+v*v
115            iu=torch.triu_indices(k,k,1,device=x.device); feat=torch.cat((u/math.sqrt(length),A[:,iu[0],iu[1]]/length,q/length),1)
116            return self.head(feat)
117
118    class GRUModel(nn.Module):
119        def __init__(self):
120            super().__init__(); self.emb=nn.Embedding(vocab,k); self.gru=nn.GRU(k,k,batch_first=True); self.head=nn.Linear(k,2)
121        def forward(self,x): return self.head(self.gru(self.emb(x))[0][:,-1])
122
123    out={}
124    for name, cls in [("additive",Additive),("levy",Levy),("gru",GRUModel)]:
125        torch.manual_seed(seed)
126        model=cls().to(device); opt=torch.optim.Adam(model.parameters(),lr=0.02); lossfn=nn.CrossEntropyLoss()
127        tx,ty=tr_x.to(device),tr_y.to(device); vx,vy=te_x.to(device),te_y.to(device)
128        t0=time.time()
129        for _ in range(epochs):
130            opt.zero_grad(); loss=lossfn(model(tx),ty); loss.backward(); opt.step()
131        with torch.no_grad():
132            pred=model(vx).argmax(1); acc=float((pred==vy).float().mean().cpu())
133        out[name]={"test_accuracy":acc,"final_train_loss":float(loss.detach().cpu()),"seconds":time.time()-t0,"parameters":sum(p.numel() for p in model.parameters())}
134    out["device"]=device
135    return out
136
137
138def benchmark(seed=123, epochs=80):
139    import torch
140    try:
141        return benchmark_impl(seed, epochs)
142    except Exception as exc:
143        # cuDNN/GPU allocation can fail on the shared device; rerun identically on CPU.
144        print("CUDA benchmark failed; falling back to CPU:", repr(exc))
145        old = torch.cuda.is_available
146        torch.cuda.is_available = lambda: False
147        try:
148            return benchmark_impl(seed, epochs)
149        finally:
150            torch.cuda.is_available = old
151
152
153if __name__ == "__main__":
154    result={"math":math_checks(),"benchmark":benchmark()}
155    with open("results.json","w") as f: json.dump(result,f,indent=2)
156    print(json.dumps(result,indent=2))