"""Regenerate only the vendored GPU shader closure from the pinned oracle."""
import subprocess,sys
from pathlib import Path
checkout=Path(sys.argv[1])
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
from mtplx.kernels import moe_sorted_gather_nax as k
assert mx.__version__=='0.32.2'
root=k._include_root();seen=set();out=[]
for rel in (k._PREAMBLE_HEADER,*k._KERNEL_HEADERS):k._inline(root,rel,seen,out)
for rel,name in [('mlx/backend/metal/kernels/unary_ops.h','Sigmoid'),('mlx/backend/metal/kernels/binary_ops.h','Multiply')]:
 functor=k._functor(root,rel,name)
 assert functor is not None
 out.append(functor)
metal=Path(__file__).resolve().parents[1]/'metal'
preamble='// MLX 0.32.2 GPU support, Copyright Apple Inc., MIT; see MLX-LICENSE.txt.\n// Flattened only the GPU header closure used by MTPLX 9882703 row-map gather.\n'
(metal/'mtplx-moe-sorted-header.metal').write_text(preamble+'\n'.join(out)+'\nusing namespace mlx::steel;\n')
(metal/'mtplx-moe-sorted-body.metal').write_text('// MTPLX 9882703, MIT; see MTPLX-LICENSE.txt.\n'+k._SOURCE)