1import sys, json, random
  2import numpy as np
  3import torch
  4from torch import nn
  5sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  6from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report
  7
  8SEEDS = tuple(range(8)); SWEEP_SEEDS = tuple(range(4))
  9NTR, NTE, EPOCHS, FT_EPOCHS = 400, 200, 22, 5
 10LRS = [1e-3, 3e-3, 1e-2]
 11KEEP = 0.50
 12
 13def seed_all(s):
 14    random.seed(s); np.random.seed(s); torch.manual_seed(s)
 15    if torch.cuda.is_available():
 16        try: torch.cuda.manual_seed_all(s)
 17        except Exception: pass
 18
 19def dataset(seed):
 20    d = get_dataset('tabular', seed=seed, n_train=NTR, n_test=NTE)
 21    # train_model expects tensors; preserve all metadata
 22    return {**d, 'xtr': torch.as_tensor(d['xtr'], dtype=torch.float32),
 23            'ytr': torch.as_tensor(d['ytr'], dtype=torch.float32),
 24            'xte': torch.as_tensor(d['xte'], dtype=torch.float32),
 25            'yte': torch.as_tensor(d['yte'], dtype=torch.float32)}
 26
 27def linear_layers(net):
 28    return [m for m in net.modules() if isinstance(m, nn.Linear)]
 29
 30def masks_for(net, x, method):
 31    ls = linear_layers(net); h = x.to(ls[0].weight.device); masks=[]
 32    for i, layer in enumerate(ls[:-1]):
 33        with torch.no_grad(): z = layer(h); h = torch.relu(z)
 34        n = z.shape[1]; k=max(1, int(round(n*KEEP)))
 35        outnorm = ls[i+1].weight.detach().cpu().numpy().T
 36        if method == 'magnitude': score=np.linalg.norm(outnorm, axis=1)
 37        else:
 38            a=z.detach().cpu().numpy()
 39            lo=np.quantile(a, .01, axis=0); hi=np.quantile(a, .99, axis=0)
 40            visible=(lo <= 0.0) & (hi >= 0.0)
 41            used=np.linalg.norm(outnorm, axis=1) > 1e-8
 42            # visibility is the structural priority; downstream norm breaks ties
 43            score=visible.astype(np.float64)*1e6 + np.linalg.norm(outnorm,axis=1)
 44            score=np.where(used, score, -1e6)
 45        idx=np.argsort(score)[-k:]
 46        m=np.zeros(n, dtype=np.float32); m[idx]=1; masks.append(m)
 47    return masks
 48
 49def apply_masks(net, masks):
 50    ls=linear_layers(net)
 51    with torch.no_grad():
 52        for i,m in enumerate(masks):
 53            t=torch.as_tensor(m, device=ls[i].weight.device)
 54            ls[i].weight.mul_(t[:,None]); ls[i].bias.mul_(t)
 55            ls[i+1].weight.mul_(t[None,:])
 56
 57def run(method, lr, seed, want_signature=False):
 58    seed_all(seed); d=dataset(seed)
 59    # This is the same base architecture and same optimizer path on both sides.
 60    net=make_model('mlp_tiny', d['input_shape'], d['out_dim'])
 61    net, _, _=train_model(net,d,epochs=EPOCHS,lr=lr,batch=128,log=lambda *a,**k:None)
 62    masks=masks_for(net,d['xtr'],method); apply_masks(net,masks)
 63    net, metric, _=train_model(net,d,epochs=FT_EPOCHS,lr=lr,batch=128,log=lambda *a,**k:None)
 64    if want_signature: return float(metric), net, d, masks
 65    return float(metric)
 66
 67def factory(method, cfg):
 68    return lambda seed: run(method, float(cfg['lr']), seed)
 69
 70def mechanism_signature():
 71    metric, net, d, masks=run('task_visible', 3e-3, 0, True)
 72    ls=linear_layers(net); x=d['xte'].to(ls[0].weight.device)
 73    predicted=[]; observed=[]
 74    # On each trained model, test the local affine/zero claim for selected axes.
 75    for i, m in enumerate(masks):
 76        with torch.no_grad():
 77            h=x
 78            for q in range(i): h=torch.relu(ls[q](h))
 79            z=ls[i](h); lo=torch.quantile(z,.01,dim=0); hi=torch.quantile(z,.99,dim=0)
 80        one=((lo>0)|(hi<0)).cpu().numpy() & (m>0)
 81        if not one.any(): continue
 82        # Compare each retained one-sided unit's ReLU contribution with its affine replacement.
 83        w=ls[i].weight.detach(); b=ls[i].bias.detach(); out=ls[i+1].weight.detach()
 84        zz=z.detach()
 85        for j in np.flatnonzero(one):
 86            pred=torch.relu(zz[:,j]); repl=zz[:,j] if float(lo[j])>0 else torch.zeros_like(zz[:,j])
 87            predicted.append(float(torch.max(torch.abs(pred-repl))))
 88            observed.append(float(torch.linalg.vector_norm(out[:,j]).cpu()))
 89    maxerr=max(predicted) if predicted else 0.0
 90    return {'prediction':'strictly one-sided selected ReLU axes are affine or zero on the test patch',
 91            'predicted_max_replacement_error':0.0, 'observed_max_replacement_error':maxerr,
 92            'tested_one_sided_axes':len(predicted), 'downstream_norm_sum':float(sum(observed)),
 93            'confirmed':bool(maxerr < 1e-5)}
 94
 95def main():
 96    grid=[{'lr':x} for x in LRS]
 97    base=sweep_baseline(lambda c: factory('magnitude',c), grid, seeds=SWEEP_SEEDS)
 98    idea_trials=[{'cfg':c,'result':evaluate(factory('task_visible',c),SEEDS)} for c in grid]
 99    best=min(idea_trials,key=lambda q:q['result']['mean'])
100    rep=make_report('tabular','mlp_tiny',base,best['result'],{
101      'track_justification':'ReLU hidden-axis pruning is structurally matched to the tabular MLP track; both systems use the identical trained mlp_tiny and optimizer budgets.',
102      'pruning_fraction':KEEP, 'idea_config':best['cfg'], 'idea_sweep':idea_trials,
103      'parameter_fraction_estimate':KEEP*KEEP, 'mechanism_signature':mechanism_signature()})
104    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
105    print(json.dumps(rep,indent=2))
106if __name__=='__main__': main()