1import json, math, random, time
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5
  6SEED = 428
  7np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  8try:
  9    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 10    if device.type == "cuda":
 11        torch.cuda.empty_cache()
 12except Exception:
 13    device = torch.device("cpu")
 14
 15
 16def orthogonal(n, seed):
 17    rng = np.random.default_rng(seed)
 18    q, r = np.linalg.qr(rng.normal(size=(n, n)))
 19    s = np.sign(np.diag(r)); s[s == 0] = 1
 20    return q @ np.diag(s)
 21
 22
 23def math_check(d=5, r=3, n_power=12):
 24    q = orthogonal(d+r, SEED)
 25    D = np.diag(np.r_[np.ones(d), -np.ones(r)])
 26    up, um = q, q @ D
 27    T, F = up[:d, :d], up[:d, d:]
 28    G, H = up[d:, :d], up[d:, d:]
 29    block_minus = np.block([[T, -F], [G, -H]])
 30    checks = {
 31        "Uplus_orth_error": float(np.linalg.norm(up.T @ up - np.eye(d+r))),
 32        "Uminus_orth_error": float(np.linalg.norm(um.T @ um - np.eye(d+r))),
 33        "block_sign_construction_error": float(np.linalg.norm(um-block_minus)),
 34        "difference_structure_error": float(np.linalg.norm((up-um)-np.block([[np.zeros((d,d)),2*F],[np.zeros((r,d)),2*H]]))),
 35        "power_difference_identity_error": 0.0,
 36        "max_norm_error_over_powers": 0.0,
 37    }
 38    diff = np.zeros_like(up)
 39    left = np.eye(d+r)
 40    right = np.eye(d+r)
 41    for k in range(n_power):
 42        # identity for N=k: U+^(k+1)-U-^(k+1)
 43        lhs = np.linalg.matrix_power(up,k+1)-np.linalg.matrix_power(um,k+1)
 44        rhs = up @ (np.linalg.matrix_power(up,k)-np.linalg.matrix_power(um,k)) + (up-um) @ np.linalg.matrix_power(um,k)
 45        checks["power_difference_identity_error"] = max(checks["power_difference_identity_error"], float(np.linalg.norm(lhs-rhs)))
 46        v = np.random.default_rng(SEED+k).normal(size=d+r)
 47        for mat in (up, um):
 48            checks["max_norm_error_over_powers"] = max(checks["max_norm_error_over_powers"], abs(np.linalg.norm(np.linalg.matrix_power(mat,k) @ v)-np.linalg.norm(v)))
 49    return checks
 50
 51class TanhRNN(nn.Module):
 52    def __init__(self, h):
 53        super().__init__(); self.h=h
 54        self.A=nn.Parameter(torch.randn(h,h)*0.12); self.B=nn.Parameter(torch.randn(h,1)*0.12); self.C=nn.Parameter(torch.randn(1,h)*0.12); self.bias=nn.Parameter(torch.zeros(h))
 55    def forward(self, x, return_norms=False):
 56        z=torch.zeros(x.size(0),self.h,device=x.device); norms=[]
 57        for t in range(x.size(1)):
 58            z=torch.tanh(z@self.A.T + x[:,t:t+1]@self.B.T + self.bias); norms.append(z.norm(dim=1).mean())
 59        return (z@self.C.T).squeeze(-1), torch.stack(norms) if return_norms else None
 60
 61class FixedUnitary(nn.Module):
 62    def __init__(self, U, h):
 63        super().__init__(); self.register_buffer('U',torch.tensor(U,dtype=torch.float32)); self.B=nn.Parameter(torch.randn(h,1)*0.12); self.C=nn.Parameter(torch.randn(1,h)*0.12)
 64    def forward(self,x,return_norms=False):
 65        z=torch.zeros(x.size(0),self.U.size(0),device=x.device); norms=[]
 66        for t in range(x.size(1)):
 67            z=z@self.U.T+x[:,t:t+1]@self.B.T; norms.append(z.norm(dim=1).mean())
 68        return (z@self.C.T).squeeze(-1), torch.stack(norms) if return_norms else None
 69
 70def make_data(batch, seq, device, generator):
 71    # The first token carries a random +/- bit; subsequent tokens are distractors.
 72    y=torch.randint(0,2,(batch,),device=device,generator=generator).float()*2-1
 73    x=torch.randn(batch,seq,1,device=device,generator=generator)*0.35
 74    x[:,0,0]=y
 75    return x,y
 76
 77def train_eval(model, seq=40, steps=350, h=32):
 78    model.to(device); opt=torch.optim.Adam(model.parameters(),lr=3e-3)
 79    g=torch.Generator(device=device); g.manual_seed(SEED+77)
 80    losses=[]
 81    for _ in range(steps):
 82        x,y=make_data(48,seq,device,g); pred,_=model(x)
 83        loss=torch.nn.functional.softplus(-pred*y).mean(); opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),10); opt.step(); losses.append(float(loss.detach().cpu()))
 84    with torch.no_grad():
 85        x,y=make_data(512,seq,device,g); pred,norms=model(x,True); acc=float(((pred>0)==(y>0)).float().mean().cpu()); test_loss=float(torch.nn.functional.softplus(-pred*y).mean().cpu()); norm_last=float(norms[-1].cpu()); norm_max=float(norms.max().cpu())
 86    # Differentiable probe is deliberately separate from no_grad evaluation.
 87    x,y=make_data(32,seq,device,g); x.requires_grad_(True); pred,_=model(x); grad=torch.autograd.grad(pred.mean(),x,retain_graph=False)[0]; gfirst=float(grad[:,0].norm(dim=1).mean().cpu()); glast=float(grad[:,-1].norm(dim=1).mean().cpu())
 88    return {"final_train_loss":float(np.mean(losses[-25:])),"test_loss":test_loss,"accuracy":acc,"hidden_norm_last":norm_last,"hidden_norm_max":norm_max,"input_grad_first":gfirst,"input_grad_last":glast}
 89
 90def main():
 91    global device
 92    h=32; q=orthogonal(h,SEED+1); d=h//2; D=np.diag(np.r_[np.ones(d),-np.ones(h-d)])
 93    t=time.time(); math_results=math_check()
 94    def run():
 95        results={"device":str(device),"math":math_results,"experiment":{}}
 96        results["experiment"]["tanh_rnn"]=train_eval(TanhRNN(h))
 97        results["experiment"]["single_unitary"]=train_eval(FixedUnitary(q,h))
 98        results["experiment"]["dual_unitary"]=train_eval(FixedUnitary(q@D,h))
 99        results["elapsed_sec"]=time.time()-t
100        return results
101    try:
102        results=run()
103    except (RuntimeError, torch.cuda.OutOfMemoryError) as exc:
104        if device.type != "cuda":
105            raise
106        print("CUDA failed; retrying on CPU:", repr(exc))
107        try: torch.cuda.empty_cache()
108        except Exception: pass
109        device=torch.device("cpu")
110        results=run()
111    with open("results.json","w") as f: json.dump(results,f,indent=2)
112    print(json.dumps(results,indent=2))
113
114if __name__=="__main__": main()