1import json, time
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6SEED = 17
  7np.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.zeros(1, device=device)
 12except Exception:
 13    device = torch.device('cpu')
 14
 15# Core recurrence from the proposal, with a compact GRU-like parameterization.
 16class EdgeGRU(nn.Module):
 17    def __init__(self, inp=2, mem=8):
 18        super().__init__(); self.mem = mem
 19        self.z = nn.Linear(inp + mem, mem); self.q = nn.Linear(inp + mem, mem)
 20        self.h = nn.Linear(inp + mem, mem)
 21    def forward(self, m, u):
 22        v = torch.cat([u, m], -1)
 23        z = torch.sigmoid(self.z(v)); q = torch.sigmoid(self.q(v))
 24        cand = torch.tanh(self.h(torch.cat([u, q*m], -1)))
 25        return (1-z)*m + z*cand
 26
 27class PersistentEdge(nn.Module):
 28    def __init__(self, mem=8):
 29        super().__init__(); self.rnn=EdgeGRU(2,mem); self.out=nn.Linear(mem,1)
 30    def forward(self, vals, active):
 31        # vals/active: [batch,T], one stable pair per sequence; dictionary is per batch key.
 32        m=torch.zeros(vals.size(0), self.rnn.mem, device=vals.device); outputs=[]
 33        for t in range(vals.size(1)):
 34            u=torch.stack([active[:,t], vals[:,t]],-1)
 35            # Inactive pairs are not updated: this is dictionary retention.
 36            new=self.rnn(m,u); m=torch.where(active[:,t,None].bool(),new,m)
 37            outputs.append(self.out(m).squeeze(-1))
 38        return torch.stack(outputs,1)
 39
 40class ActiveMLP(nn.Module):
 41    def __init__(self):
 42        super().__init__(); self.net=nn.Sequential(nn.Linear(2,16),nn.Tanh(),nn.Linear(16,1))
 43    def forward(self, vals, active):
 44        return self.net(torch.stack([active,vals],-1)).squeeze(-1)
 45
 46def math_check():
 47    torch.manual_seed(3)
 48    g=EdgeGRU(2,16); m=torch.randn(1000,16)*20; u=torch.randn(1000,2)*50
 49    with torch.no_grad():
 50        mp=g(m,u)
 51        # Since candidate is in [-1,1], each coordinate is a convex combination.
 52        bound=(mp.abs() <= torch.maximum(m.abs(), torch.ones_like(m))).float().mean().item()
 53        inactive=torch.randn(4,16); update=g(inactive,torch.randn(4,2))
 54        retained=torch.where(torch.zeros(4,1).bool(),update,inactive)
 55        exact=(retained==inactive).all().item()
 56    return {'convex_coordinate_bound_fraction':bound,'inactive_retention_exact':bool(exact),
 57            'max_abs_before':float(m.abs().max()),'max_abs_after':float(mp.abs().max())}
 58
 59def data(n):
 60    # Pair-specific latent contact history. At t=0..2 the value is observable;
 61    # after a gap it reappears at t=9..11 with no informative current value.
 62    latent=torch.randint(0,2,(n,)).float()*2-1
 63    T=12; active=torch.zeros(n,T); vals=torch.zeros(n,T)
 64    active[:,0:3]=1; active[:,9:12]=1; vals[:,0:3]=latent[:,None]
 65    target=latent[:,None].expand(n,T)
 66    return vals.to(device),active.to(device),target.to(device)
 67
 68def train(model, vals, active, target, steps=700):
 69    opt=torch.optim.Adam(model.parameters(),lr=3e-3)
 70    model.train()
 71    for _ in range(steps):
 72        # fresh batch each step, avoiding memorization of sequence ordering
 73        v,a,y=data(128); pred=model(v,a)
 74        loss=((pred[:,9:12]-y[:,9:12])**2).mean()
 75        opt.zero_grad(); loss.backward(); opt.step()
 76    model.eval()
 77    with torch.no_grad():
 78        pred=model(vals,active); react=((pred[:,9:12]-target[:,9:12])**2).mean().item()
 79        one=((pred[:,0:3]-target[:,0:3])**2).mean().item()
 80        long=((pred-target)**2).mean().item()
 81    return react,one,long
 82
 83def timed(model, vals, active, reps=100):
 84    with torch.no_grad():
 85        for _ in range(10): model(vals,active)
 86        if device.type=='cuda': torch.cuda.synchronize()
 87        t=time.perf_counter()
 88        for _ in range(reps): model(vals,active)
 89        if device.type=='cuda': torch.cuda.synchronize()
 90    return (time.perf_counter()-t)/reps*1000
 91
 92def main():
 93    check=math_check(); vals,active,target=data(2048)
 94    torch.manual_seed(SEED); base=ActiveMLP().to(device)
 95    b=train(base,vals,active,target)
 96    torch.manual_seed(SEED); idea=PersistentEdge(8).to(device)
 97    e=train(idea,vals,active,target)
 98    base_ms=timed(base,vals[:128],active[:128]); edge_ms=timed(idea,vals[:128],active[:128])
 99    result={'device':str(device),'math_check':check,
100      'baseline':{'reactivation_mse':b[0],'early_mse':b[1],'all_timestep_mse':b[2], 'params':sum(p.numel() for p in base.parameters()),'ms_batch128':base_ms},
101      'persistent_edge_gru':{'reactivation_mse':e[0],'early_mse':e[1],'all_timestep_mse':e[2], 'params':sum(p.numel() for p in idea.parameters()),'ms_batch128':edge_ms},
102      'reactivation_mse_ratio_idea_over_baseline':e[0]/max(b[0],1e-12),
103      'latency_overhead_ratio':edge_ms/max(base_ms,1e-12)}
104    print(json.dumps(result,indent=2))
105if __name__=='__main__': main()