1import json, math, random, time
 2import numpy as np
 3import torch
 4from torch import nn
 5
 6SEED=2275
 7random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
 8device='cuda' if torch.cuda.is_available() else 'cpu'
 9try:
10    if device=='cuda': torch.cuda.empty_cache()
11except Exception:
12    device='cpu'
13
14# Three mechanism checks, directly tied to the proposed update and halting equations.
15def mechanism_checks():
16    alpha=0.08; eps=1e-3; n=30
17    # For s_(t+1)=(1+alpha*lambda)s_t, boundedness is |1+alpha*lambda|<=1,
18    # with strict contraction for <1. The two neutral boundaries are lambda=0,-2/alpha.
19    stability=[]
20    for lam in [-26,-25,-20,-13,-12.5,-10,0,5]:
21        q=1+alpha*lam; vals=np.array([q**t for t in range(n+1)])
22        pred='contract' if abs(q)<1 else ('neutral' if abs(q)==1 else 'diverge')
23        obs='diverge' if max(abs(vals))>1+1e-8 else ('neutral' if abs(q)==1 else 'contract')
24        stability.append({'lambda':lam,'q':q,'predicted':pred,'observed':obs,'abs_s30':float(abs(vals[-1]))})
25    # Formula T=min{t: sum_{k=1}^t h_k >= 1-eps}; for constant h this is ceil((1-eps)/h).
26    hrows=[]
27    for h in [.05,.1,.2,.4,.7]:
28        pred=math.ceil((1-eps)/h)
29        cumulative=0.; t=0
30        while cumulative < 1-eps:
31            t+=1; cumulative+=h
32        # ACT clipping reaches mass exactly one on the same final update, but its
33        # stopping criterion is mass after clipping; report it separately.
34        mass=0.; act_t=0
35        while mass < 1-eps:
36            act_t+=1; mass+=min(h,1-mass)
37        hrows.append({'h':h,'formula_T':pred,'observed_formula_T':t,'act_clipped_T':act_t})
38    # Since h=sigmoid(b), the predicted T must be non-increasing in b.
39    brows=[]
40    for b in [-3,-1,0,1,3]:
41        h=1/(1+math.exp(-b)); pred=math.ceil((1-eps)/h)
42        cumulative=0.; t=0
43        while cumulative < 1-eps: t+=1; cumulative+=h
44        brows.append({'b':b,'h':h,'formula_T':pred,'observed_T':t})
45    return {'stability':stability,'strict_stability_interval':[-25.0,0.0],
46            'halting_constant_h':hrows,'halting_bias_sweep':brows}
47
48class ResidualMLP(nn.Module):
49    def __init__(self,d=32,blocks=6):
50        super().__init__(); self.inp=nn.Linear(2,d); self.blocks=nn.ModuleList([nn.Sequential(nn.LayerNorm(d),nn.Linear(d,d),nn.GELU(),nn.Linear(d,d)) for _ in range(blocks)]); self.out=nn.Linear(d,2)
51    def forward(self,x):
52        s=self.inp(x)
53        for b in self.blocks: s=s+b(s)
54        return self.out(s)
55
56class AdaptiveRecurrence(nn.Module):
57    def __init__(self,d=32,tmax=8):
58        super().__init__(); self.inp=nn.Linear(2,d); self.norm=nn.LayerNorm(d); self.rule=nn.Sequential(nn.Linear(d,48),nn.GELU(),nn.Linear(48,d)); self.halt=nn.Linear(d,1); self.out=nn.Linear(d,2); self.tmax=tmax
59    def forward(self,x):
60        s=self.inp(x); acc=torch.zeros_like(s); mass=torch.zeros(x.shape[0],1,device=x.device); steps=torch.zeros_like(mass)
61        for _ in range(self.tmax):
62            s=s+0.25*self.rule(self.norm(s)); h=torch.sigmoid(self.halt(s)); delta=torch.minimum(h,1-mass); acc=acc+delta*s; mass=mass+delta; steps=steps+(mass<1-1e-3).float()
63        acc=acc+(1-mass)*s
64        return self.out(acc), steps.squeeze(1)
65
66def make_data(n, seed):
67    g=np.random.default_rng(seed); x=g.uniform(-1,1,(n,2)).astype('float32'); y=((x[:,0]**2+x[:,1]**2 + .22*x[:,0])>.42).astype('int64')
68    return torch.tensor(x),torch.tensor(y)
69
70def train(model, steps=500):
71    model.to(device); x,y=make_data(2048,SEED+1); xv,yv=make_data(2048,SEED+2); x,y,xv,yv=[z.to(device) for z in (x,y,xv,yv)]
72    opt=torch.optim.Adam(model.parameters(),lr=3e-3); lossfn=nn.CrossEntropyLoss(); t0=time.perf_counter(); model.train()
73    for i in range(steps):
74        ix=torch.randint(0,len(x),(128,),device=device); xb,yb=x[ix],y[ix]; opt.zero_grad(); out=model(xb); logits=out[0] if isinstance(out,tuple) else out; loss=lossfn(logits,yb)
75        if isinstance(out,tuple): loss=loss+0.002*out[1].mean()
76        loss.backward(); opt.step()
77    if device=='cuda': torch.cuda.synchronize()
78    elapsed=time.perf_counter()-t0; model.eval()
79    with torch.no_grad():
80        out=model(xv); logits=out[0] if isinstance(out,tuple) else out; acc=(logits.argmax(1)==yv).float().mean().item(); avgsteps=out[1].mean().item() if isinstance(out,tuple) else 6.0
81    return {'accuracy':acc,'avg_microsteps':avgsteps,'train_seconds':elapsed,'parameters':sum(p.numel() for p in model.parameters())}
82
83def main():
84    result={'device':device,'mechanism_checks':mechanism_checks(),'mini_experiment':{'baseline':train(ResidualMLP()),'idea':train(AdaptiveRecurrence())}}
85    with open('results.json','w') as f: json.dump(result,f,indent=2)
86    print(json.dumps(result,indent=2))
87if __name__=='__main__': main()