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

torch.manual_seed(0)
class MinState(nn.Module):
    def __init__(self):
        super().__init__()
        self.register_buffer("bank", torch.zeros(1, 64, 256))  # all dims 32-friendly
        self.enc = nn.Linear(256, 256)
    def forward(self, x):  # (1,64,256)
        y = self.enc(x + self.bank)
        self.bank.copy_(y)  # whole-buffer update, no slice/cat
        return y

m = MinState().eval(); m.requires_grad_(False)
x = torch.randn(1, 64, 256)
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()
o = ml.predict({"x": x.numpy()}, state=st)
print("minimal aligned whole-buffer state on CPU_AND_NE: predict OK")