1import json, math, os, time
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6SEED = 7
  7torch.manual_seed(SEED); np.random.seed(SEED)
  8try:
  9    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 10    if device.type == 'cuda':
 11        torch.zeros(1, device=device)
 12except Exception:
 13    device = torch.device('cpu')
 14
 15# A small input-conditioned linear recurrent layer: h' = A(z)h + g(z).
 16class ICRec(nn.Module):
 17    def __init__(self, d=8, k=3, mode='free', rho=.90, lam=3.0):
 18        super().__init__(); self.d=d; self.k=k; self.mode=mode; self.rho=rho; self.lam=lam
 19        self.A = nn.Parameter(.25*torch.randn(k,d,d))
 20        self.inp = nn.Linear(1,d)
 21        self.gate = nn.Linear(1,k)
 22        self.out = nn.Linear(d,1)
 23        # P = L L^T + eps I, with positive diagonal through softplus.
 24        self.rawL = nn.Parameter(torch.eye(d) + .03*torch.randn(d,d))
 25        self.log_diag = nn.Parameter(torch.zeros(d))
 26
 27    def P(self):
 28        L = torch.tril(self.rawL, diagonal=-1) + torch.diag(torch.nn.functional.softplus(self.log_diag)+.08)
 29        return L @ L.T + 1e-4*torch.eye(self.d, device=L.device)
 30
 31    def matrices(self):
 32        if self.mode == 'spectral':
 33            # Standard pointwise Euclidean spectral normalization, applied separately.
 34            vals=[]
 35            for a in self.A:
 36                s=torch.linalg.svdvals(a)[0]
 37                vals.append(a * (.90 / (s + 1e-6)).clamp(max=1.0))
 38            return torch.stack(vals)
 39        return self.A
 40
 41    def lyap_penalty(self):
 42        P=self.P(); mats=self.matrices()
 43        # Generalized eigenvalues of A^T P A v = lambda P v.
 44        # Cholesky transform is symmetric and differentiable for this small MVP.
 45        L=torch.linalg.cholesky(P)
 46        invL=torch.linalg.inv(L)
 47        vals=[]
 48        for a in mats:
 49            B=invL @ (a.T @ P @ a) @ invL.T
 50            vals.append(torch.linalg.eigvalsh((B+B.T)/2)[-1])
 51        vmax=torch.stack(vals)
 52        violation=torch.relu(vmax-self.rho**2)
 53        return self.lam*(violation**2).sum(), vmax.detach(), violation.detach()
 54
 55    def forward(self,x, return_stats=False):
 56        h=torch.zeros(x.shape[0],self.d,device=x.device); hs=[]
 57        mats=self.matrices()
 58        for t in range(x.shape[1]):
 59            z=x[:,t,:]
 60            pi=torch.softmax(self.gate(z),-1)
 61            A=torch.einsum('bk,kij->bij',pi,mats)
 62            h=torch.bmm(A,h.unsqueeze(-1)).squeeze(-1)+self.inp(z)
 63            h=torch.tanh(h)
 64            hs.append(h)
 65        y=self.out(h).squeeze(-1)
 66        if return_stats: return y, torch.stack(hs,1)
 67        return y
 68
 69def math_check():
 70    # Random positive P and matrices deliberately rescaled to satisfy the CQLF bound.
 71    torch.manual_seed(SEED+1); d=5; k=4; rho=.82
 72    L=torch.randn(d,d); P=L@L.T+.5*torch.eye(d); C=torch.linalg.cholesky(P)
 73    mats=[]
 74    for _ in range(k):
 75        a=torch.randn(d,d)
 76        # scale in P geometry so max generalized eigenvalue is below rho^2.
 77        invC=torch.linalg.inv(C)
 78        B=invC @ (a.T@P@a) @ invC.T
 79        lm=torch.linalg.eigvalsh((B+B.T)/2)[-1]
 80        mats.append(a*(rho*.72/torch.sqrt(lm)))
 81    mats=torch.stack(mats)
 82    # Every arbitrary convex gate and every length-25 product must contract in P norm.
 83    max_violation=-1.; max_ratio=0.; worst=0.
 84    for trial in range(200):
 85        h=torch.randn(d); e0=h@P@h
 86        for t in range(25):
 87            pi=torch.softmax(torch.randn(k),0); a=torch.einsum("k,kij->ij", pi, mats); h=a@h
 88            ratio=torch.sqrt((h@P@h)/e0).item()
 89            max_ratio=max(max_ratio, ratio/(rho**(t+1)))
 90        # ratio/rho^T <= 1 is the claimed bound
 91    for a in mats:
 92        invC=torch.linalg.inv(C); B=invC@(a.T@P@a)@invC.T
 93        max_violation=max(max_violation,(torch.linalg.eigvalsh((B+B.T)/2)[-1]-rho**2).item())
 94    return {'max_basis_violation':max_violation, 'max_normalized_product_ratio':max_ratio,
 95            'rho':rho, 'pass': bool(max_violation <= 2e-5 and max_ratio <= 1.0001)}
 96
 97def train(mode, steps=350):
 98    torch.manual_seed(SEED); np.random.seed(SEED)
 99    m=ICRec(mode=mode).to(device); opt=torch.optim.Adam(m.parameters(),lr=3e-3)
100    t0=time.time(); losses=[]; final_h=0.; final_grad=0.; penalty=0.
101    for step in range(steps):
102        # First symbol is the bit to remember; remaining symbols are zero.
103        x=torch.zeros(96,30,1,device=device); bit=torch.randint(0,2,(96,),device=device).float()
104        x[:,0,0]=bit
105        target=2*bit-1
106        pred,h=m(x,True); task=nn.functional.mse_loss(pred,target)
107        p,_,_=m.lyap_penalty() if mode=='lyap' else (torch.tensor(0.,device=device),None,None)
108        loss=task+p; opt.zero_grad(); loss.backward()
109        final_grad=float(torch.nn.utils.clip_grad_norm_(m.parameters(),100.).detach().cpu()); opt.step()
110        losses.append(float(task.detach().cpu())); final_h=float(h.detach().abs().max().cpu()); penalty=float(p.detach().cpu())
111    # Diagnostic random long products using learned basis matrices.
112    with torch.no_grad():
113        mats=m.matrices(); P=m.P(); C=torch.linalg.cholesky(P); invC=torch.linalg.inv(C)
114        eig=[]
115        for a in mats:
116            B=invC@(a.T@P@a)@invC.T; eig.append(float(torch.linalg.eigvalsh((B+B.T)/2)[-1].cpu())**.5)
117        x=torch.zeros(256,60,1,device=device); bit=torch.randint(0,2,(256,),device=device).float(); x[:,0,0]=bit
118        pred,h=m(x,True); acc=float(((pred>0)==(bit>0)).float().mean().cpu())
119    return {'loss_start':losses[0], 'loss_end':losses[-1], 'accuracy_long':acc,
120            'max_hidden_last_batch':final_h, 'last_grad_norm':final_grad,
121            'lyap_penalty':penalty, 'max_basis_P_gain':max(eig), 'seconds':time.time()-t0}
122
123def main():
124    check=math_check(); results={}
125    for mode in ['free','spectral','lyap']:
126        results[mode]=train(mode)
127    out={'device':str(device), 'math_check':check, 'results':results}
128    print(json.dumps(out,indent=2))
129
130if __name__=='__main__': main()