File size: 871 Bytes
47d8ad6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 | 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")
|