1import sys, json, random
 2from pathlib import Path
 3import numpy as np
 4import torch
 5from torch import nn
 6
 7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
 8from bench import get_dataset, train_model, evaluate, sweep_baseline, make_report, reload_custom_tracks
 9
10HERE = Path(__file__).parent
11reload_custom_tracks()
12TRACK = 'correlated_token_moe_regression'
13SEEDS = tuple(range(8))
14# Shared union: all idea learning rates are also baseline sweep values.
15GRID = [{'lr': 1e-3, 'epochs': 20}, {'lr': 3e-3, 'epochs': 20}, {'lr': 1e-2, 'epochs': 20}]
16
17class Baseline(nn.Module):
18    """Standard label-sensitive flattened MLP, matched width/depth."""
19    def __init__(self):
20        super().__init__()
21        self.net = nn.Sequential(nn.Linear(64, 64), nn.Tanh(),
22                                 nn.Linear(64, 64), nn.Tanh(),
23                                 nn.Linear(64, 1))
24    def forward(self, x):
25        return self.net(x.reshape(x.shape[0], -1))
26
27class Equivariant(nn.Module):
28    """Permutation-equivariant token encoder followed by invariant pooling."""
29    def __init__(self):
30        super().__init__()
31        self.token = nn.Sequential(nn.Linear(4, 64), nn.Tanh(), nn.Linear(64, 64), nn.Tanh())
32        self.pair = nn.Sequential(nn.Linear(4, 32), nn.Tanh(), nn.Linear(32, 16), nn.Tanh())
33        self.readout = nn.Sequential(nn.Linear(4 + 64 + 16, 64), nn.Tanh(),
34                                     nn.Linear(64, 64), nn.Tanh(), nn.Linear(64, 1))
35    def forward(self, x):
36        h = self.token(x)
37        mean_x = x.mean(dim=1)
38        dif = x[:, :, None, :] - x[:, None, :, :]
39        messages = self.pair(dif).mean(dim=2).mean(dim=1)
40        return self.readout(torch.cat([mean_x, h.mean(dim=1), messages], dim=1))
41
42def seed_all(seed):
43    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
44
45def run(kind, cfg, seed, keep=False):
46    seed_all(seed)
47    ds = get_dataset(TRACK, seed, 400, 200)
48    model = Baseline() if kind == 'baseline' else Equivariant()
49    net, metric, _ = train_model(model, ds, epochs=cfg['epochs'], lr=cfg['lr'], batch=128)
50    if net is None:
51        raise RuntimeError('benchmark training failed')
52    return (float(metric), net, ds) if keep else float(metric)
53
54def factory(kind):
55    return lambda cfg: (lambda seed: run(kind, cfg, seed))
56
57def main():
58    baseline = sweep_baseline(factory('baseline'), GRID)
59    idea_configs = []
60    for cfg in GRID:
61        res = evaluate(lambda seed, cfg=cfg: run('idea', cfg, seed), SEEDS)
62        idea_configs.append((cfg, res))
63    best_cfg, idea = min(idea_configs, key=lambda z: z[1]['mean'])
64
65    base_swap, idea_swap = [], []
66    for seed in SEEDS:
67        _, bnet, ds = run('baseline', baseline['best_cfg'], seed, True)
68        _, inet, _ = run('idea', best_cfg, seed, True)
69        x = ds['xte'][:64]
70        xp = x[:, [1, 0] + list(range(2, x.shape[1]))]
71        bd = next(bnet.parameters()).device; idv = next(inet.parameters()).device
72        with torch.no_grad():
73            base_swap.append(float((bnet(x.to(bd)) - bnet(xp.to(bd))).abs().mean()))
74            idea_swap.append(float((inet(x.to(idv)) - inet(xp.to(idv))).abs().mean()))
75    signature = {
76        'quantity': 'mean absolute prediction change under swapping two trained-task tokens',
77        'prediction': 'equivariant/invariant system has zero change while flattened baseline is nonzero',
78        'baseline_per_seed': base_swap,
79        'idea_per_seed': idea_swap,
80        'baseline_mean': float(np.mean(base_swap)),
81        'idea_mean': float(np.mean(idea_swap)),
82        'confirmed': bool(np.mean(idea_swap) < 1e-6 and np.mean(base_swap) > 1e-5)
83    }
84    report = make_report(TRACK, 'matched_custom', baseline, idea, signature)
85    report['idea_sweep'] = [{'cfg': cfg, 'mean': res['mean'], 'std': res['std'], 'per_seed': res['per_seed']} for cfg, res in idea_configs]
86    report['track_justification'] = 'The registered correlated-token track contains exchangeable multi-token correlations; built-in tracks do not.'
87    (HERE / 'bench_report.json').write_text(json.dumps(report, indent=2))
88    print(json.dumps(report, indent=2))
89
90if __name__ == '__main__':
91    main()