1import os, sys, json, math, random
  2import numpy as np
  3import torch
  4import torch.nn as 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))
  9# n=8 tangent subspace keeps the compound penalty small enough for the bench.
 10N_TAN=8
 11HIDDEN=64
 12
 13def seed_all(s):
 14    random.seed(s); np.random.seed(s); torch.manual_seed(s)
 15    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)
 16
 17def pairs(n): return [(i,j) for i in range(n) for j in range(i+1,n)]
 18
 19def additive_compound(A):
 20    n=A.shape[-1]; ps=pairs(n); m=len(ps)
 21    B=A.new_zeros(A.shape[:-2]+(m,m))
 22    for a,(i,j) in enumerate(ps):
 23        B[...,a,a]=A[...,i,i]+A[...,j,j]
 24        for b,(k,l) in enumerate(ps):
 25            if a != b:
 26                # coefficient of E_ij in A^[2] acting on e_k wedge e_l
 27                B[...,a,b] = (A[...,i,k]*(1.0 if j==l else 0.0)
 28                              + A[...,j,l]*(1.0 if i==k else 0.0)
 29                              - A[...,i,l]*(1.0 if j==k else 0.0)
 30                              - A[...,j,k]*(1.0 if i==l else 0.0))
 31    return B
 32
 33def wedge_math_check():
 34    torch.manual_seed(7); A=torch.randn(5,5,dtype=torch.float64)
 35    # finite difference derivative of multiplicative compound, using minors
 36    ps=pairs(5)
 37    def wedge(M):
 38        W=torch.empty(10,10,dtype=M.dtype)
 39        for a,(i,j) in enumerate(ps):
 40            for b,(k,l) in enumerate(ps): W[a,b]=M[i,k]*M[j,l]-M[i,l]*M[j,k]
 41        return W
 42    hs=torch.tensor([.02,.01,.005,.0025],dtype=torch.float64); errs=[]
 43    for h in hs: errs.append(float(torch.linalg.norm(wedge(torch.matrix_exp(h*A))-torch.eye(10)-h*additive_compound(A))))
 44    slope=float(np.polyfit(np.log(hs.numpy()),np.log(errs),1)[0])
 45    return {'identity_error_slope':slope,'expected':2.0,'max_error':max(errs),'passed':abs(slope-2)<.12}
 46
 47def cell8(model, x3, h8):
 48    """First eight rows/columns of the exact GRU Jacobian, analytically."""
 49    W=model.rnn.weight_ih_l0; U=model.rnn.weight_hh_l0
 50    b1=model.rnn.bias_ih_l0; b2=model.rnn.bias_hh_l0
 51    h=torch.cat([h8, torch.zeros(h8.shape[0], HIDDEN-N_TAN, device=h8.device, dtype=h8.dtype)], dim=1)
 52    gi=torch.nn.functional.linear(x3,W,b1); gh=torch.nn.functional.linear(h,U,b2)
 53    ar=gi[:,:HIDDEN]+gh[:,:HIDDEN]; aq=gi[:,HIDDEN:2*HIDDEN]+gh[:,HIDDEN:2*HIDDEN]
 54    r=torch.sigmoid(ar); q=torch.sigmoid(aq)
 55    an=gi[:,2*HIDDEN:3*HIDDEN]+r*gh[:,2*HIDDEN:3*HIDDEN]
 56    n=torch.tanh(an); d=gh[:,2*HIDDEN:3*HIDDEN]
 57    # Batch is one in tangent_penalty; retain differentiability wrt parameters.
 58    dr=r*(1-r)
 59    dq=q*(1-q)
 60    # Each derivative tensor has shape [1,HIDDEN,N_TAN].
 61    Ur=U[:HIDDEN,:N_TAN].unsqueeze(0)
 62    Uq=U[HIDDEN:2*HIDDEN,:N_TAN].unsqueeze(0)
 63    Un=U[2*HIDDEN:3*HIDDEN,:N_TAN].unsqueeze(0)
 64    drdh=dr.unsqueeze(2)*Ur
 65    dqdh=dq.unsqueeze(2)*Uq
 66    dndh=(1-n*n).unsqueeze(2)*(r.unsqueeze(2)*Un+d.unsqueeze(2)*drdh)
 67    dy=(1-q).unsqueeze(2)*dndh + dqdh*(h[:,:HIDDEN].unsqueeze(2)-n.unsqueeze(2))
 68    dy=dy + q.unsqueeze(2)*torch.eye(HIDDEN,device=h.device,dtype=h.dtype)[:,:N_TAN].unsqueeze(0)
 69    return ((1-q)*n+q*h[:,:HIDDEN])[:,:N_TAN], dy[:,:N_TAN,:]
 70
 71def tangent_penalty(model, x):
 72    x3=x[:, -3:][:1]
 73    h=torch.zeros(1,N_TAN,device=x.device)
 74    _, A=cell8(model,x3,h)
 75    A=A[0]
 76    C=additive_compound(A)
 77    off=C-torch.diag(torch.diagonal(C))
 78    neg=torch.relu(-off)
 79    return (neg*neg).mean()+0.02*torch.relu(0.01-off).pow(2).mean(), A.detach(), C.detach()
 80
 81def train_idea(seed,cfg, return_sig=False):
 82    seed_all(seed); ds=get_dataset('dynamics',seed,n_train=200,n_test=100)
 83    model=make_model('rnn_small',ds['input_shape'],ds['out_dim'])
 84    device='cuda' if torch.cuda.is_available() else 'cpu'
 85    try:
 86        model.to(device); opt=torch.optim.Adam(model.parameters(),lr=cfg['lr'],weight_decay=cfg['wd'])
 87        lossf=nn.MSELoss(); xtr,ytr=ds['xtr'].to(device),ds['ytr'].to(device)
 88        for ep in range(cfg['epochs']):
 89            model.train(); perm=torch.randperm(len(xtr),device=device)
 90            for i in range(0,len(xtr),64):
 91                ix=perm[i:i+64]; pred=model(xtr[ix]); task=lossf(pred,ytr[ix])
 92                # Applying at every mini-batch makes the intervention part of training.
 93                pen,_,_=tangent_penalty(model,xtr[ix])
 94                loss=task+cfg['lam']*pen
 95                opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),5); opt.step()
 96        model.eval()
 97        with torch.no_grad(): metric=float(((model(ds['xte'].to(device))-ds['yte'].to(device))**2).mean())
 98        if return_sig:
 99            with torch.enable_grad(): pen,A,C=tangent_penalty(model,xtr[:1])
100            off=C-torch.diag(torch.diagonal(C)); return metric, {'negative_offdiag_fraction':float((off<0).float().mean()),'mean_negative_magnitude':float(torch.relu(-off).mean()),'penalty':float(pen),'jacobian_norm':float(torch.linalg.norm(A))}
101        return metric
102    except RuntimeError:
103        # CPU fallback for tight CUDA/cuDNN environments.
104        seed_all(seed); os.environ['CUDA_VISIBLE_DEVICES']=''
105        return train_idea_cpu(seed,cfg,return_sig)
106
107def train_idea_cpu(seed,cfg,return_sig=False):
108    # same implementation after disabling CUDA; recursive body is intentionally avoided.
109    old=torch.cuda.is_available
110    ds=get_dataset('dynamics',seed,n_train=200,n_test=100); model=make_model('rnn_small',ds['input_shape'],ds['out_dim'])
111    opt=torch.optim.Adam(model.parameters(),lr=cfg['lr'],weight_decay=cfg['wd']); lossf=nn.MSELoss(); xtr,ytr=ds['xtr'],ds['ytr']
112    for _ in range(cfg['epochs']):
113        for i in range(0,len(xtr),64):
114            task=lossf(model(xtr[i:i+64]),ytr[i:i+64]); pen,_,_=tangent_penalty(model,xtr[i:i+64]); opt.zero_grad(); (task+cfg['lam']*pen).backward(); opt.step()
115    with torch.no_grad(): metric=float(((model(ds['xte'])-ds['yte'])**2).mean())
116    if return_sig:
117        pen,A,C=tangent_penalty(model,xtr[:1]); off=C-torch.diag(torch.diagonal(C)); return metric,{'negative_offdiag_fraction':float((off<0).float().mean()),'mean_negative_magnitude':float(torch.relu(-off).mean()),'penalty':float(pen),'jacobian_norm':float(torch.linalg.norm(A))}
118    return metric
119
120def base_fn(cfg):
121    def run(seed):
122        seed_all(seed); ds=get_dataset('dynamics',seed,n_train=200,n_test=100); m=make_model('rnn_small',ds['input_shape'],ds['out_dim'])
123        _,metric,_=train_model(m,ds,epochs=cfg['epochs'],lr=cfg['lr'],batch=64,weight_decay=cfg['wd'],log=lambda *_:None)
124        return metric
125    return run
126
127def main():
128    check=wedge_math_check(); lrs=[0.0015,0.003,0.006]; wds=[0.0]
129    grid=[{'lr':lr,'wd':wd,'epochs':3} for lr in lrs for wd in [0.0]]
130    base=sweep_baseline(base_fn,grid,seeds=(0,1,2,3))
131    # Same lr/weight-decay union; lambda sweep is the sole method difference.
132    ideas=[]
133    for lr in lrs:
134        for lam in [0.001,0.005,0.02]: ideas.append({'lr':lr,'wd':base['best_cfg']['wd'],'epochs':3,'lam':lam})
135    best=None; rows=[]
136    for cfg in ideas:
137        r=evaluate(lambda s,cfg=cfg: train_idea(s,cfg),seeds=SEEDS); rows.append({'cfg':cfg,'mean':r['mean']})
138        if best is None or r['mean']<best['mean']: best=dict(r); best['cfg']=cfg
139    sigvals=[train_idea(s,best['cfg'],True) for s in SEEDS]; best['per_seed']=[v[0] for v in sigvals]; best['mean']=float(np.mean(best['per_seed'])); best['std']=float(np.std(best['per_seed'])); best['n']=8
140    sig={'predicted':'Metzler violation should decrease in trained recurrent tangent dynamics; lower violation should accompany positive-cone mechanism.','observed_mean':{k:float(np.mean([v[1][k] for v in sigvals])) for k in sigvals[0][1]},'confirmed':False}
141    report=make_report('dynamics','rnn_small',base,best,{'math_sanity':check,'idea_sweep':rows,**sig})
142    with open('bench_report.json','w') as f: json.dump(report,f,indent=2)
143    print(json.dumps(report,indent=2))
144if __name__=='__main__': main()