"""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]))