import numpy as np, torch, torch.nn as nn, sys
import coremltools as ct

C, S, T = 256, 64, 7
torch.manual_seed(0)

class Stateless(nn.Module):
    def __init__(self):
        super().__init__()
        self.attn = nn.MultiheadAttention(C, 4, batch_first=True)
        self.enc = nn.Linear(C, C)
    def forward(self, x):
        y, _ = self.attn(x, x, x)
        return self.enc(y)

class StateNoAttn(nn.Module):
    def __init__(self):
        super().__init__()
        self.register_buffer("bank", torch.zeros(T, S, C))
        self.enc = nn.Linear(C, C)
    def forward(self, x):
        y = self.enc(x) + self.bank.mean(dim=0, keepdim=True)
        self.bank.copy_(torch.cat([self.bank[1:], y], 0))
        return y

which = sys.argv[1]
m = {"stateless": Stateless, "state_noattn": StateNoAttn}[which]().eval()
m.requires_grad_(False)
x = torch.randn(1, S, C)
with torch.no_grad():
    ep = torch.export.export(m, (x,)); ep = ep.run_decompositions({})
ml = ct.convert(ep, minimum_deployment_target=ct.target.iOS18,
                compute_units=ct.ComputeUnit.CPU_AND_NE)
st = ml.make_state() if which != "stateless" else None
kw = {"state": st} if st is not None else {}
o = ml.predict({"x": x.numpy()}, **kw)
print(which, "on CPU_AND_NE: predict OK")