1import json, sys, random, importlib.util, math
 2from pathlib import Path
 3import numpy as np
 4import torch
 5import torch.nn as nn
 6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
 7from bench import make_model, train_model, evaluate, sweep_baseline, make_report
 8
 9OUT=Path('bench_report.json'); SEEDS=tuple(range(8)); SWEEP=(0,1,2,3)
10LR_GRID=[1e-3,3e-3,1e-2]; EPOCHS=12; NTR=400; NTE=200
11spec=importlib.util.spec_from_file_location('local_track','/home/maxwelhelp/all/math2nn/bench/custom_tracks/conditional_multitoken_diffusion.py')
12track=importlib.util.module_from_spec(spec); spec.loader.exec_module(track)
13
14def seed_all(s):
15    random.seed(s); np.random.seed(s); torch.manual_seed(s)
16    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)
17
18def ds(seed):
19    d0=track.get_dataset(seed,NTR,NTE)
20    return {'xtr':torch.as_tensor(d0['xtr'],dtype=torch.float32),'ytr':torch.as_tensor(d0['ytr'],dtype=torch.float32),
21            'xte':torch.as_tensor(d0['xte'],dtype=torch.float32),'yte':torch.as_tensor(d0['yte'],dtype=torch.float32),
22            'task':'regression','metric':'mse','input_shape':tuple(d0['xtr'].shape[1:]),'out_dim':8}
23
24class MixtureNet(nn.Module):
25    def __init__(self, shape, k=3):
26        super().__init__(); self.base=make_model('transformer_tiny',shape,16*k); self.k=k
27    def forward(self,x):
28        q=self.base(x); return q[:,:8*self.k].reshape(-1,8,self.k), q[:,8*self.k:].reshape(-1,8,self.k)
29
30def idea_train(seed,lr,return_sig=False):
31    seed_all(seed); d=ds(seed); device='cuda' if torch.cuda.is_available() else 'cpu'
32    try: net=MixtureNet(d['input_shape']).to(device)
33    except Exception: device='cpu'; net=MixtureNet(d['input_shape']).to(device)
34    opt=torch.optim.Adam(net.parameters(),lr=lr); x,y=d['xtr'].to(device),d['ytr'].to(device)
35    for _ in range(EPOCHS):
36        for ix in torch.randperm(len(x),device=device).split(128):
37            mu,lg=net(x[ix]); yy=y[ix].unsqueeze(-1); sd=.12
38            lp=torch.log_softmax(lg,2)-.5*((yy-mu)/sd)**2-math.log(sd)-.5*math.log(2*math.pi)
39            loss=-torch.logsumexp(lp,2).mean(); opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(net.parameters(),5); opt.step()
40    net.eval()
41    with torch.no_grad():
42        mu,lg=net(d['xte'].to(device)); pred=(torch.softmax(lg,2)*mu).sum(2); yy=d['yte'].to(device)
43        mse=((pred-yy)**2).mean().item(); resid=(yy-pred).cpu().numpy(); v=resid.var()
44        kurt=float(np.mean((resid-resid.mean())**4)/(v*v+1e-12)-3); spread=float((mu.max(2).values-mu.min(2).values).mean())
45    return (mse,{'mean_component_spread':spread,'residual_excess_kurtosis':kurt}) if return_sig else mse
46
47def base_train(seed,lr):
48    seed_all(seed); d=ds(seed); net=make_model('transformer_tiny',d['input_shape'],d['out_dim'])
49    return train_model(net,d,epochs=EPOCHS,lr=lr,batch=128)[1]
50
51def main():
52    # Normal PIT sanity check for the inverse-CDF transport construction.
53    rng=np.random.default_rng(1342); z=rng.normal(size=20000); u=0.5*(1+np.vectorize(math.erf)(z/math.sqrt(2)))
54    pit=float(np.max(np.abs(np.sort(u)-(np.arange(len(u))+.5)/len(u))))
55    grid=[{'lr':v} for v in LR_GRID]
56    base=sweep_baseline(lambda c:lambda s:base_train(s,c['lr']),grid,seeds=SWEEP)
57    tried=[]
58    for lr in LR_GRID: tried.append(evaluate(lambda s,lr=lr:idea_train(s,lr),seeds=SEEDS)|{'cfg':{'lr':lr,'components':3}})
59    best=min(tried,key=lambda r:r['mean']); sigs=[idea_train(s,best['cfg']['lr'],True)[1] for s in SEEDS]
60    sig={k:float(np.mean([z[k] for z in sigs])) for k in sigs[0]}; sig.update({'prediction':'multimodal conditional transport has separated learned components','confirmed':sig['mean_component_spread']>0.05})
61    rep=make_report('conditional_multitoken_diffusion','transformer_tiny',base,best,extra=sig)
62    rep['idea_sweep']=tried; rep['math_check']={'normal_inverse_cdf_pit_sup_deviation':pit,'expected_below':0.02}; rep['custom_track']={'name':'conditional_multitoken_diffusion','file':'/home/maxwelhelp/all/math2nn/bench/custom_tracks/conditional_multitoken_diffusion.py','domain':'diffusion-sampling'}
63    rep['protocol_notes']='Matched eight-token diffusion task; independently trained shared transformer backbone; identical test MSE readout; baseline and idea share lr union.'
64    OUT.write_text(json.dumps(rep,indent=2)); print(json.dumps(rep,indent=2))
65if __name__=='__main__': main()