1import json
 2import numpy as np
 3import torch
 4from torch import nn
 5
 6SEED=965
 7
 8def pk_loss(z,x,eps=1e-6):
 9    u=z/(torch.linalg.vector_norm(z,dim=-1,keepdim=True)+eps)
10    r2=(x*x).sum().clamp(max=1-eps)
11    den=1+r2-2*(u*x).sum(-1)
12    return (-torch.log((1-r2)/den.clamp_min(eps))).mean()
13
14class RNN(nn.Module):
15    def __init__(self):
16        super().__init__(); self.W=nn.Linear(3,2); self.U=nn.Linear(2,2,bias=False); self.out=nn.Linear(2,1)
17    def forward(self,s):
18        h=torch.zeros(s.shape[0],2,device=s.device); hs=[]; ys=[]
19        for t in range(s.shape[1]):
20            h=torch.tanh(self.W(torch.cat([s[:,t],h],-1))); hs.append(h); ys.append(self.out(h))
21        return torch.stack(ys,1),torch.stack(hs,1)
22
23def run(reg):
24    torch.manual_seed(SEED); np.random.seed(SEED)
25    dev='cuda' if torch.cuda.is_available() else 'cpu'
26    try:
27        device=torch.device(dev); model=RNN().to(device)
28        g=torch.Generator(device=device); g.manual_seed(SEED)
29        ntr,nva,T=192,64,25
30        t=torch.arange(ntr+nva+T+1,device=device).float()
31        base=torch.sin(.18*t)+.15*torch.sin(.73*t)
32        noise=.08*torch.randn(ntr+nva+T+1,generator=g,device=device)
33        seq=(base+noise).unfold(0,T+1,1)[:ntr+nva]
34        x=seq[:,:T].unsqueeze(-1); y=seq[:,1:T+1].unsqueeze(-1)
35        opt=torch.optim.Adam(model.parameters(),lr=.015)
36        losses=[]
37        # fixed affine random-map composition estimate, as the implementation plan specifies
38        q=.72; target=torch.tensor([.48,.20],device=device); probes=torch.tensor([[-.5,.1],[.1,-.4],[.4,.2]],device=device)
39        xhat=target+(q**8)*(probes.mean(0)-target)
40        for step in range(180):
41            idx=torch.randperm(ntr,generator=g,device=device)[:64]
42            pred,h=model(x[idx]); task=((pred-y[idx])**2).mean(); loss=task
43            if reg: loss=loss+.01*pk_loss(h[:,-1],xhat.detach())
44            opt.zero_grad(); loss.backward(); opt.step(); losses.append(float(task.detach()))
45        with torch.no_grad():
46            pred,h=model(x[ntr:]); mse=float(((pred-y[ntr:])**2).mean())
47            z=h[:,-1]; pair=torch.pdist(z).mean().item(); pl=pk_loss(z,xhat).item()
48        return {'val_mse':mse,'pairwise_distance':pair,'pk_nll':pl,'last_train_task':losses[-1]}
49    except Exception as e:
50        if dev=='cuda':
51            torch.cuda.empty_cache(); return run_cpu(reg)
52        raise
53
54def run_cpu(reg):
55    old=torch.cuda.is_available
56    torch.cuda.is_available=lambda:False
57    try:return run(reg)
58    finally:torch.cuda.is_available=lambda:old
59
60if __name__=='__main__':
61    out={'baseline':run(False),'poisson_kernel':run(True)}
62    with open('mini_results.json','w') as f: json.dump(out,f,indent=2)
63    print(json.dumps(out,indent=2))