20 lines
1.2 KiB
Python
20 lines
1.2 KiB
Python
"""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)
|