| """Does ANY stateful attention model ANE-compile on this machine? | |
| Reuses the cross-state toy from test_cross_state.py at small scale.""" | |
| import numpy as np, torch, sys | |
| import coremltools as ct | |
| from test_cross_state import CrossStateModel, C, S | |
| torch.manual_seed(0) | |
| m = CrossStateModel().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({}) | |
| unit = getattr(ct.ComputeUnit, sys.argv[1] if len(sys.argv) > 1 else "CPU_AND_NE") | |
| ml = ct.convert(ep, minimum_deployment_target=ct.target.iOS18, compute_units=unit) | |
| st = ml.make_state() | |
| o1 = ml.predict({"x": x.numpy()}, state=st); o2 = ml.predict({"x": x.numpy()}, state=st) | |
| k = list(o1)[0] | |
| print("toy stateful attn on", unit, "OK; state evolves:", not np.allclose(o1[k], o2[k])) | |