1import sys, json, random, heapq
 2from pathlib import Path
 3import numpy as np
 4import torch
 5import torch.nn.functional as F
 6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
 7from bench import get_dataset, make_model, sweep_baseline, make_report, evaluate
 8
 9SEEDS=tuple(range(8)); GRID=[{'lr':1e-3,'epochs':12},{'lr':3e-3,'epochs':12},{'lr':6e-3,'epochs':12}]; BATCH=128; CHUNKS=8
10class Node:
11    def __init__(self,a=None,b=None,leaf=None): self.a,self.b,self.leaf=a,b,leaf
12def balanced(ids):
13    ids=list(ids)
14    if len(ids)==1:return Node(leaf=ids[0])
15    m=len(ids)//2; return Node(balanced(ids[:m]),balanced(ids[m:]))
16def huffman(w):
17    h=[(float(x),i,Node(leaf=i)) for i,x in enumerate(w)]; heapq.heapify(h); serial=len(h)
18    while len(h)>1:
19        a,_,x=heapq.heappop(h); b,_,y=heapq.heappop(h); heapq.heappush(h,(a+b,serial,Node(x,y))); serial+=1
20    return h[0][2]
21def depths(t,d=0,o=None):
22    o={} if o is None else o
23    if t.leaf is not None:o[t.leaf]=d
24    else:depths(t.a,d+1,o);depths(t.b,d+1,o)
25    return o
26def seed_all(s):
27    random.seed(s);np.random.seed(s);torch.manual_seed(s)
28def round_stages(x,n):
29    # Shape-preserving proxy for n internal low-precision additions. Adding zero
30    # isolates rounding while avoiding the invalid operation of summing disjoint chunks.
31    z=torch.zeros_like(x)
32    for _ in range(n): x=(x.to(torch.float16)+z.to(torch.float16)).to(torch.float32)
33    return x
34def run(seed,cfg,adaptive,signature=False):
35    seed_all(seed); ds=get_dataset('tabular',seed=seed,n_train=400,n_test=400)
36    net=make_model('mlp_tiny',ds['input_shape'],ds['out_dim']); opt=torch.optim.AdamW(net.parameters(),lr=cfg['lr'])
37    xtr,ytr=ds['xtr'],ds['ytr']; rng=np.random.default_rng(seed+991); em=np.ones(CHUNKS); eq=np.ones(CHUNKS); last=None; errors=[]
38    net.train()
39    for _ in range(cfg['epochs']):
40        perm=rng.permutation(len(xtr))
41        for st in range(0,len(xtr),BATCH):
42            idx=torch.as_tensor(perm[st:st+BATCH]); loss=F.mse_loss(net(xtr[idx]),ytr[idx]); opt.zero_grad();loss.backward()
43            ps=[p for p in net.parameters() if p.grad is not None]; flat=torch.cat([p.grad.detach().reshape(-1) for p in ps]); raw=list(torch.tensor_split(flat,CHUNKS))
44            means=np.array([float(x.mean()) for x in raw]); secs=np.array([float((x*x).mean()) for x in raw]); em=.9*em+.1*means;eq=.9*eq+.1*secs;var=np.maximum(eq-em*em,1e-12)
45            tree=huffman(var) if adaptive else balanced(range(CHUNKS)); last=depths(tree)
46            approx=torch.cat([round_stages(x,last[i]) for i,x in enumerate(raw)]); errors.append(float(((approx-flat)**2).mean()))
47            pos=0
48            for p in ps:n=p.numel();p.grad.copy_(approx[pos:pos+n].reshape_as(p));pos+=n
49            opt.step()
50    with torch.no_grad(): metric=float(((net(ds['xte'])-ds['yte'])**2).mean())
51    if signature:return metric,{'last_depths':{str(k):int(v) for k,v in last.items()},'observed_reduction_mse':float(np.mean(errors)),'n_steps':len(errors)}
52    return metric
53def baseline_fn(c):return lambda s:run(s,c,False)
54def idea_fn(c):return lambda s:run(s,c,True)
55def main():
56    base=sweep_baseline(baseline_fn,GRID); tried=[]
57    for c in GRID:tried.append({'cfg':c,'mean':float(np.mean([idea_fn(c)(s) for s in range(4)]))})
58    best=min(tried,key=lambda x:x['mean'])['cfg']; idea=evaluate(idea_fn(best),seeds=SEEDS)
59    ss=[run(s,best,True,True)[1] for s in SEEDS]; da=[x['last_depths'] for x in ss]; md=float(np.mean([min(map(int,d.values())) for d in da]))
60    sig={'predicted':'high-variance chunks have shallower Huffman depth','observed_mean_min_depth':md,'balanced_depth':3,'trained_model_measurements':da,'observed_reduction_mse_mean':float(np.mean([x['observed_reduction_mse'] for x in ss])),'confirmed':bool(md<3)}
61    rep=make_report('tabular','mlp_tiny',base,idea,{'idea_sweep':tried,'mechanism_signature':sig});rep['idea']['best_cfg']=best;rep['idea']['sweep']=tried;Path('bench_report.json').write_text(json.dumps(rep,indent=2));print(json.dumps(rep,indent=2))
62if __name__=='__main__':main()