"""Pinned head-chunk plan, actual SDPA subcall shapes and exact BF16 outputs."""
import hashlib,io,json,os,subprocess,sys
from pathlib import Path
checkout,out=map(Path,sys.argv[1:3])
assert subprocess.check_output(['git','-C',str(checkout),'rev-parse','HEAD'],text=True).strip()=='9882703f3105363ddc37eca9f97aa09a1d387112'
sys.path.insert(0,str(checkout))
import mlx.core as mx
import numpy as np
from mtplx.models import qwen4_exp as qwen
assert mx.__version__=='0.32.2'
plans=[dict(rows=r,heads=h,kv=kv,plan=qwen._verify_sdpa_head_chunk_plan(r,h,kv)) for r in [-1,0,1,2,3,4,8,9,16,31,32,33,64] for h in [0,1,2,12,24,33,36,64] for kv in [0,1,2,3,5]]
def pattern(shape,salt):
 n=int(np.prod(shape));v=((np.arange(n,dtype=np.uint32)*17+salt*13)%257).astype(np.int32)-128
 return mx.array((v.astype(np.float32)/128).reshape(shape)).astype(mx.bfloat16)
geometries=[(1,24,2,r,65) for r in [1,2,3,4,8,9,16,32,33]]+[(2,24,2,4,17),(1,64,1,1,17),(1,14,2,5,17),(1,12,2,5,17),(1,1,1,4,17),(1,24,2,4,2),(1,24,2,4,1024),(1,24,2,4,4097),(1,5,2,4,17)]
original=mx.fast.scaled_dot_product_attention
calls=[]
def traced(q,k,v,**kwargs):
 calls.append(dict(q=list(q.shape),k=list(k.shape),v=list(v.shape),mask=None if kwargs['mask'] is None else list(kwargs['mask'].shape)))
 return original(q,k,v,**kwargs)
mx.fast.scaled_dot_product_attention=traced
cases=[]
for batch,heads,kv,rows,total in geometries:
 for mode in ['none','bool','add','heads','local','two_dim']:
  for enabled in [False,True]:
   os.environ['MTPLX_QWEN4_VERIFY_SDPA_HEAD_CHUNK']='1' if enabled else '0'
   calls.clear();case=dict(batch=batch,heads=heads,kv=kv,rows=rows,total=total,mode=mode,enabled=enabled)
   q=pattern((batch,rows,heads,256),93).transpose(0,2,1,3)
   k=pattern((batch,kv,total+7,256),94)[:,:,2:total+2]
   v=pattern((batch,kv,total+7,256),95)[:,:,2:total+2]
   mask=None
   if mode!='none':
    mh=heads if mode=='heads' else (4 if mode=='local' else 1)
    mask=np.array([j<=total-rows+i and (j+h)%3!=1 and i!=1 for h in range(mh) for i in range(rows) for j in range(total)],dtype=np.bool_).reshape(1,mh,rows,total)
    mask=mx.array(mask)
    if mode=='add':mask=mx.where(mask,mx.array(0,dtype=mx.bfloat16),mx.array(-10,dtype=mx.bfloat16))
    if mode=='two_dim':mask=mask.reshape(rows,total)
   try:
    y=qwen._verify_sdpa(q,k,v,scale=.0625,mask=mask)
    dot=io.StringIO();mx.export_to_dot(dot,output=y)
    case['fused_primitives']=dot.getvalue().count('ScaledDotProductAttention')
    mx.eval(y)
    case['sha256']=hashlib.sha256(np.asarray(y.view(mx.uint16)).tobytes()).hexdigest()
    case['shape']=list(y.shape)
   except ValueError as e:case['error']=str(e)
   case['calls']=list(calls);cases.append(case)
   del q,k,v,mask
   mx.clear_cache()
assert len(cases)==216
out.write_text(json.dumps(dict(revision='9882703f3105363ddc37eca9f97aa09a1d387112',mlx=mx.__version__,plans=plans,cases=cases),separators=(',',':'))+'\n')
print(json.dumps(dict(plans=len(plans),cases=len(cases),errors=sum('error' in c for c in cases),output=str(out))))