1import sys, json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report
  8
  9SEED0 = 2230
 10EPOCHS = 12
 11NTR, NTE = 800, 300
 12BATCH = 128
 13
 14
 15def sinkhorn(logits, tau=1.0, steps=10):
 16    # Stable log-domain alternating normalization; output is nonnegative DS.
 17    z = logits / tau
 18    for _ in range(steps):
 19        z = z - torch.logsumexp(z, dim=-1, keepdim=True)
 20        z = z - torch.logsumexp(z, dim=-2, keepdim=True)
 21    return torch.exp(z)
 22
 23
 24class StreamTransformer(nn.Module):
 25    """S streams of d features, with shared transformer-style feature processing."""
 26    def __init__(self, input_dim, out_dim, S=4, ds=16, depth=2, tau=1.0, sink_steps=10, alpha=.1, constrained=True):
 27        super().__init__()
 28        self.S, self.ds, self.tau, self.sink_steps, self.constrained = S, ds, tau, sink_steps, constrained
 29        self.inp = nn.Linear(1, S * ds)
 30        self.pos = nn.Parameter(torch.zeros(1, 32, S, ds))
 31        self.g = nn.Parameter(torch.eye(S).unsqueeze(0).repeat(depth, 1, 1) * 2.0)
 32        self.blocks = nn.ModuleList()
 33        for _ in range(depth):
 34            self.blocks.append(nn.ModuleDict({
 35                'norm': nn.LayerNorm(ds),
 36                'ff': nn.Sequential(nn.Linear(ds, 64), nn.GELU(), nn.Linear(64, ds))
 37            }))
 38        self.alpha = nn.Parameter(torch.tensor(float(alpha)))
 39        self.out = nn.Linear(32 * S * ds, out_dim)
 40
 41    def forward(self, x):
 42        X = self.inp(x.unsqueeze(-1)).view(x.shape[0], x.shape[1], self.S, self.ds)
 43        X = X + self.pos[:, :x.shape[1]]
 44        for i, block in enumerate(self.blocks):
 45            A = sinkhorn(self.g[i], self.tau, self.sink_steps) if self.constrained else self.g[i]
 46            Xm = torch.einsum('ij,btjd->btid', A, X)
 47            # Shared streamwise residual transformation, preserving stream routing.
 48            F = block['ff'](block['norm'](Xm))
 49            X = Xm + self.alpha * F
 50        return self.out(X.reshape(x.shape[0], -1))
 51
 52
 53def set_seed(s):
 54    random.seed(s); np.random.seed(s); torch.manual_seed(s)
 55    if torch.cuda.is_available():
 56        try: torch.cuda.manual_seed_all(s)
 57        except Exception: pass
 58
 59
 60def train_one(kind, seed, lr, tau=1.0, alpha=.1, return_model=False):
 61    set_seed(seed)
 62    d = get_dataset('sequence', seed, n_train=NTR, n_test=NTE)
 63    if kind == 'baseline':
 64        model = StreamTransformer(d['input_shape'][0], d['out_dim'], tau=tau, alpha=alpha, constrained=False)
 65    else:
 66        model = StreamTransformer(d['input_shape'][0], d['out_dim'], tau=tau, alpha=alpha)
 67    net, metric, hist = train_model(model, d, epochs=EPOCHS, lr=lr, batch=BATCH, log=lambda *_: None)
 68    if net is None: return (float('inf'), None, d) if return_model else float('inf')
 69    if return_model: return float(metric), net, d
 70    return float(metric)
 71
 72
 73def signature(seed=0, lr=3e-3):
 74    metric, net, d = train_one('idea', seed, lr, return_model=True)
 75    device = next(net.parameters()).device
 76    with torch.no_grad():
 77        xt = d['xte'][:64].to(device)
 78        X = net.inp(xt.unsqueeze(-1)).view(xt.shape[0], 32, net.S, net.ds)
 79        ratios, re, ce = [], [], []
 80        for g in net.g:
 81            A = sinkhorn(g, net.tau, net.sink_steps)
 82            Xm = torch.einsum('ij,btjd->btid', A, X)
 83            ratios.append(float(torch.linalg.norm(Xm) / torch.linalg.norm(X)))
 84            re.append(float((A.sum(-1)-1).abs().max())); ce.append(float((A.sum(-2)-1).abs().max()))
 85    # Prediction tested on trained model: DS mixing is non-expansive (allow tiny FP tolerance).
 86    return {'trained_model': True, 'predicted': {'max_frobenius_ratio': '<=1', 'row_col_error': '0'},
 87            'observed': {'max_frobenius_ratio': max(ratios), 'mean_frobenius_ratio': float(np.mean(ratios)),
 88                         'max_row_error': max(re), 'max_col_error': max(ce)},
 89            'confirmed': bool(max(ratios) <= 1.0001 and max(re) < 1e-4 and max(ce) < 1e-4),
 90            'metric_on_signature_model': float(metric)}
 91
 92
 93def main():
 94    # Union parity: every idea learning rate is also evaluated by baseline.
 95    grid = [{'lr': 7e-4}, {'lr': 1e-3}, {'lr': 1.5e-3}, {'lr': 2e-3}, {'lr': 3e-3}]
 96    base = sweep_baseline(lambda c: lambda s: train_one('baseline', s, c['lr']), grid)
 97    best_lr = base['best_cfg']['lr']
 98    idea_grid = [1e-3, 1.5e-3, 2e-3]
 99    idea_runs = []
100    for lr in idea_grid:
101        r = evaluate(lambda s, lr=lr: train_one('idea', s, lr))
102        idea_runs.append({'lr': lr, 'result': r})
103    best = min(idea_runs, key=lambda z: z['result']['mean'])
104    rep = make_report('sequence', 'transformer_tiny', base, best['result'], {
105        'idea_lr_sweep': idea_runs, 'selected_lr': best['lr'], **signature(0, best['lr'])})
106    rep['custom_track'] = None
107    Path('bench_report.json').write_text(json.dumps(rep, indent=2))
108    print(json.dumps(rep, indent=2))
109
110if __name__ == '__main__': main()