1import os, sys, json, random
  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, train_model, sweep_baseline, make_report
  8
  9SEED0 = 597
 10LR_GRID = [1e-3, 3e-3, 1e-2]
 11EPOCHS = 12
 12BATCH = 128
 13
 14
 15def canonical(word, rank, supports):
 16    a = list(word)
 17    changed = True
 18    while changed:
 19        changed = False
 20        for i in range(len(a)-1):
 21            x, y = a[i], a[i+1]
 22            commute = supports[x].isdisjoint(supports[y])
 23            if commute and rank[x] > rank[y]:
 24                a[i], a[i+1] = y, x
 25                changed = True
 26    return tuple(a)
 27
 28
 29class ModularSequence(nn.Module):
 30    """Tiny sequence predictor with trace-canonical modular sub-block."""
 31    def __init__(self, word):
 32        super().__init__()
 33        self.width = 32
 34        self.inp = nn.Linear(1, self.width)
 35        self.pos = nn.Parameter(torch.zeros(1, 32, self.width))
 36        # Operators have disjoint read/write supports for g0..g3.
 37        self.ops = nn.ModuleDict()
 38        self.supports = {}
 39        for g in range(4):
 40            name = f"g{g}"
 41            self.ops[name] = nn.Sequential(nn.Linear(8, 8), nn.GELU())
 42            self.supports[name] = frozenset(range(8*g, 8*(g+1)))
 43        self.ops["global"] = nn.Sequential(nn.Linear(32, 32), nn.GELU())
 44        self.supports["global"] = frozenset(range(32))
 45        self.order = tuple(word)
 46        self.rank = {n:i for i,n in enumerate(("g0","g1","g2","g3","global"))}
 47        self.head = nn.Sequential(nn.LayerNorm(32), nn.Linear(32*32, 1))
 48
 49    def apply_op(self, h, name):
 50        if name == "global":
 51            return self.ops[name](h)
 52        idx = sorted(self.supports[name])
 53        z = self.ops[name](h[..., idx])
 54        out = h.clone()
 55        out[..., idx] = z
 56        return out
 57
 58    def forward_with_order(self, x, word):
 59        h = self.inp(x.unsqueeze(-1)) + self.pos[:, :x.shape[1]]
 60        for name in word:
 61            h = self.apply_op(h, name)
 62        h = self.head[0](h)
 63        return self.head[1](h.reshape(h.shape[0], -1))
 64
 65    def forward(self, x):
 66        return self.forward_with_order(x, self.order)
 67
 68
 69def make_net(seed, idea):
 70    torch.manual_seed(seed)
 71    # Deliberately noncanonical proposal; only disjoint group swaps are legal.
 72    proposed = ("g3", "g1", "global", "g0", "g2")
 73    word = canonical(proposed, {n:i for i,n in enumerate(("g0","g1","g2","g3","global"))},
 74                     {**{f"g{i}": frozenset(range(8*i,8*(i+1))) for i in range(4)}, "global":frozenset(range(32))}) if idea else proposed
 75    return ModularSequence(word)
 76
 77
 78def train_one(seed, lr, idea, return_model=False):
 79    # Identical task, model family, optimizer, epochs, and batch size.
 80    torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)
 81    d = get_dataset("sequence", seed, n_train=4000, n_test=1000)
 82    net = make_net(seed, idea)
 83    trained, metric, hist = train_model(net, d, epochs=EPOCHS, lr=lr, batch=BATCH, log=lambda *_: None)
 84    if trained is None:
 85        raise RuntimeError("bench training failed")
 86    return (float(metric), trained, d) if return_model else float(metric)
 87
 88
 89def signature():
 90    vals_comm, vals_non = [], []
 91    for seed in range(8):
 92        _, net, d = train_one(seed, 3e-3, True, True)
 93        x = d["xte"][:64].to(next(net.parameters()).device)
 94        prop = net.order
 95        can = canonical(prop, net.rank, net.supports)
 96        with torch.no_grad():
 97            vals_comm.append(float((net.forward_with_order(x, prop)-net.forward_with_order(x, can)).abs().max()))
 98            # Swap a group operator across global: this is not support-disjoint.
 99            bad = ("g3", "g1", "g0", "global", "g2")
100            vals_non.append(float((net.forward_with_order(x, prop)-net.forward_with_order(x, bad)).abs().max()))
101    mc, mn = max(vals_comm), float(np.median(vals_non))
102    return {"prediction": "disjoint-support swaps preserve the trained network output; group/global swaps do not",
103            "predicted_commuting_max": 1e-5, "observed_commuting_max": mc,
104            "predicted_noncommuting_positive": True, "observed_noncommuting_median": mn,
105            "confirmed": bool(mc <= 1e-5 and mn > 1e-4), "n_trained_models": 8}
106
107
108def main():
109    # Baseline sweep includes every idea-side learning rate (search-space parity).
110    grid = [{"lr": lr, "epochs": EPOCHS, "batch": BATCH} for lr in LR_GRID]
111    def factory(cfg):
112        return lambda seed: train_one(seed, cfg["lr"], False)
113    base = sweep_baseline(factory, grid)
114    # Idea sweep: best baseline lr plus two nearby/equal-budget settings.
115    idea_runs = []
116    for lr in LR_GRID:
117        r = {"lr": lr, "epochs": EPOCHS, "batch": BATCH,
118             "result": {"per_seed": [train_one(s, lr, True) for s in range(8)]}}
119        r["result"]["mean"] = float(np.mean(r["result"]["per_seed"]))
120        r["result"]["std"] = float(np.std(r["result"]["per_seed"]))
121        r["result"]["n"] = 8
122        idea_runs.append(r)
123    best = min(idea_runs, key=lambda z:z["result"]["mean"])
124    report = make_report("sequence", "modular_sequence_local", base, best["result"],
125                         {"mechanism_signature": signature(),
126                          "track_justification": "sequence is structurally matched: multi-token window correlations and a sequence-level model",
127                          "idea_sweep": [{"cfg": {k:v for k,v in r.items() if k != "result"}, "mean": r["result"]["mean"]} for r in idea_runs],
128                          "canonical_proposal": ["g3","g1","global","g0","g2"],
129                          "canonical_word": list(canonical(("g3","g1","global","g0","g2"), {n:i for i,n in enumerate(("g0","g1","g2","g3","global"))}, {**{f"g{i}":frozenset(range(8*i,8*(i+1))) for i in range(4)},"global":frozenset(range(32))}))})
130    with open("bench_report.json", "w") as f: json.dump(report, f, indent=2)
131    print(json.dumps(report, indent=2))
132
133if __name__ == "__main__": main()