1import json, random, numpy as np, torch
 2from experiment import AttnRegressor, make_data
 3
 4def one(seed, device):
 5    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 6    tr_x,tr_y=make_data(256); te_x,te_y=make_data(64)
 7    tr_x,tr_y,te_x,te_y=[z.to(device) for z in (tr_x,tr_y,te_x,te_y)]
 8    out={}
 9    for biased in (False,True):
10        torch.manual_seed(seed); m=AttnRegressor(biased).to(device)
11        opt=torch.optim.Adam(m.parameters(),lr=3e-3); mse=torch.nn.MSELoss()
12        for _ in range(180):
13            ix=torch.randint(0,256,(16,),device=device)
14            loss=mse(m(tr_x[ix]),tr_y[ix]); opt.zero_grad(); loss.backward(); opt.step()
15        with torch.no_grad():
16            pred=m(te_x); err=(pred-te_y).square().squeeze(-1)
17            r=(te_x.square().sum(-1)).sqrt()
18            near=err[r<.35].mean().item(); far=err[r>.7].mean().item()
19            out['idea' if biased else 'baseline']=[mse(pred,te_y).item(),near,far]
20    return out
21
22def main():
23    device='cuda' if torch.cuda.is_available() else 'cpu'
24    try:
25        vals=[one(s,device) for s in (17,23,41)]
26    except Exception:
27        device='cpu'; vals=[one(s,device) for s in (17,23,41)]
28    print(json.dumps({'device':device,'seeds':[17,23,41],'columns':['overall_mse','near_origin_mse','far_mse'],'runs':vals},indent=2))
29if __name__=='__main__': main()