1import os, sys, json, math
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7from bench import get_dataset, make_model, evaluate, sweep_baseline, make_report
  8
  9SEED0 = 1616
 10BATCH = 64
 11EPOCHS = 18
 12TAUS = [0.0, 0.05, 0.10]
 13LRS = [1e-3, 3e-3, 6e-3]
 14BETAS = [0.8, 0.9]
 15
 16
 17def polar(x):
 18    # Batched-free, exact polar factor for a 2-D weight/gradient matrix.
 19    u, _, vh = torch.linalg.svd(x, full_matrices=False)
 20    return u @ vh
 21
 22
 23def matrix_sign(x):
 24    return torch.where(x >= 0, torch.ones_like(x), -torch.ones_like(x))
 25
 26
 27def direction(m, residual, tau, heldout, gated):
 28    post = matrix_sign(polar(m))
 29    pre = polar(matrix_sign(m + residual))
 30    rho = torch.sum(heldout * post) / (torch.linalg.vector_norm(heldout) *
 31                                         torch.linalg.vector_norm(post) + 1e-8)
 32    if gated and float(rho.detach()) < tau:
 33        return pre, post, pre, float(rho.detach()), False
 34    return post, post, pre, float(rho.detach()), True
 35
 36
 37def train_one(seed, lr, beta, method, tau=0.05, collect=False):
 38    # Explicit CPU fallback also avoids inheriting a failed CUDA context.
 39    torch.manual_seed(seed); np.random.seed(seed)
 40    ds = get_dataset('tabular', seed, n_train=400, n_test=200)
 41    device = 'cuda' if torch.cuda.is_available() else 'cpu'
 42    try:
 43        net = make_model('mlp_tiny', ds['input_shape'], ds['out_dim']).to(device)
 44    except Exception:
 45        device = 'cpu'
 46        net = make_model('mlp_tiny', ds['input_shape'], ds['out_dim']).to(device)
 47    x, y = ds['xtr'].to(device), ds['ytr'].to(device)
 48    xt, yt = ds['xte'].to(device), ds['yte'].to(device)
 49    lossf = nn.MSELoss()
 50    # SignMuon-like state only for 2-D weight matrices; vectors use Adam-style SGD.
 51    mom = {p: torch.zeros_like(p) for p in net.parameters() if p.ndim == 2}
 52    residual = {p: torch.zeros_like(p) for p in net.parameters() if p.ndim == 2}
 53    rng = np.random.default_rng(seed + 991)
 54    logs = {'post_fraction': [], 'rho': [], 'post_negative': [], 'chosen_negative': [], 'predicted_fraction': []}
 55    net.train()
 56    n = x.shape[0]
 57    for epoch in range(EPOCHS):
 58        order = rng.permutation(n)
 59        # The next cyclic minibatch is an independent routing estimate.
 60        for bi in range(0, n, BATCH):
 61            ids = order[bi:bi+BATCH]
 62            hid = order[(bi+BATCH) % n:(bi+2*BATCH) % n] if bi+BATCH < n else order[:min(BATCH,n)]
 63            # avoid accidental empty/wrapped slices
 64            if len(hid) == 0: hid = order[:BATCH]
 65            net.zero_grad(set_to_none=True)
 66            pred = net(x[ids].view(len(ids), -1))
 67            lossf(pred, y[ids]).backward()
 68            # Save current gradients, then independently estimate heldout gradients.
 69            grads = {p: (p.grad.detach().clone() if p.grad is not None else torch.zeros_like(p))
 70                     for p in net.parameters()}
 71            net.zero_grad(set_to_none=True)
 72            lossf(net(x[hid].view(len(hid), -1)), y[hid]).backward()
 73            held = {p: (p.grad.detach().clone() if p.grad is not None else torch.zeros_like(p))
 74                    for p in net.parameters()}
 75            with torch.no_grad():
 76                for p in net.parameters():
 77                    g = grads[p]
 78                    if p.ndim == 2:
 79                        mom[p].mul_(beta).add_(g, alpha=1-beta)
 80                        r = residual[p]
 81                        d, post, pre, rho, use_post = direction(mom[p], r, tau, held[p], method == 'switch')
 82                        p.add_(d, alpha=-lr / math.sqrt(max(1, p.shape[1])))
 83                        if method == 'switch' and not use_post:
 84                            residual[p].copy_(mom[p] + r - matrix_sign(mom[p] + r))
 85                        elif method == 'pre':
 86                            residual[p].copy_(mom[p] + r - matrix_sign(mom[p] + r))
 87                        if collect:
 88                            logs['post_fraction'].append(float(use_post))
 89                            logs['rho'].append(rho)
 90                            logs['post_negative'].append(float(torch.sum(held[p]*post).item() < 0))
 91                            logs['chosen_negative'].append(float(torch.sum(held[p]*d).item() < 0))
 92                            logs['predicted_fraction'].append(float(rho >= tau))
 93                    else:
 94                        # Same simple gradient-side update for all non-matrix parameters.
 95                        p.add_(g, alpha=-lr)
 96    net.eval()
 97    with torch.no_grad():
 98        metric = float(lossf(net(xt.view(len(xt), -1)), yt).cpu())
 99    if collect:
100        logs = {k: float(np.mean(v)) if v else 0.0 for k,v in logs.items()}
101        logs['n_observations'] = int(len(order) * EPOCHS / BATCH * max(1, len(mom)))
102    return metric, logs
103
104
105def fn(cfg, method):
106    return lambda seed: train_one(seed, cfg['lr'], cfg['beta'], method, cfg.get('tau', 0.05))[0]
107
108
109def main():
110    # Union parity: every idea lr/beta is included in baseline's grid.
111    grid = [{'lr': lr, 'beta': beta} for lr in LRS for beta in BETAS]
112    base = sweep_baseline(lambda c: fn(c, 'post'), grid)
113    idea_grid = [dict(c, tau=t) for c in grid for t in TAUS]
114    # Evaluate the idea's 3 tau choices at the selected baseline settings and two nearby lr settings.
115    bcfg = base['best_cfg']; nearby = sorted(set([bcfg['lr']] + [v for v in LRS if v != bcfg['lr']]))[:3]
116    candidates = [{'lr': lr, 'beta': bcfg['beta'], 'tau': tau} for lr in nearby for tau in TAUS]
117    tried = []
118    for c in candidates:
119        r = evaluate(fn(c, 'switch'))
120        tried.append((c, r))
121    best_cfg, idea = min(tried, key=lambda z: z[1]['mean'])
122    sigs = [train_one(s, best_cfg['lr'], best_cfg['beta'], 'switch', best_cfg['tau'], True)[1] for s in range(8)]
123    sig = {k: float(np.mean([z[k] for z in sigs])) for k in sigs[0] if k != 'n_observations'}
124    sig['confirmed'] = bool(abs(sig['post_fraction'] - sig['predicted_fraction']) < 1e-9 and sig['chosen_negative'] <= sig['post_negative'])
125    base['sweep'] = base['sweep']
126    report = make_report('tabular', 'mlp_tiny', base, idea, {
127        'prediction': 'heldout alignment gating routes post when rho>=tau and reduces negative heldout alignment',
128        'observed': sig, 'tau': best_cfg['tau'], 'trained_models': True
129    })
130    report['idea_sweep'] = [{'cfg': c, 'mean': r['mean'], 'per_seed': r['per_seed']} for c,r in tried]
131    report['custom_track'] = None
132    with open('bench_report.json','w') as f: json.dump(report, f, indent=2)
133    print(json.dumps(report, indent=2))
134
135if __name__ == '__main__': main()