mlx: add compiled closure support

Wraps MLX's mlx_compile API so Go functions can be traced into fused
kernels. Contiguous elementwise chains collapse into a single
Metal/CUDA kernel instead of launching one per op.

Exposes Compile plus arity helpers (Compile1/2/3) that mirror Python's
@mx.compile decorator shape, lazily building the closure on first call
so package-level declarations work before the MLX dylib loads.
This commit is contained in:
Jesse Gross
2026-04-13 12:20:33 -07:00
parent 698e04a14b
commit d3e67e305c
6 changed files with 374 additions and 30 deletions

View File

@@ -27,7 +27,11 @@ var arrays []*Array
func New(name string) *Array {
t := &Array{name: name}
arrays = append(arrays, t)
if tracing {
traceScratch = append(traceScratch, t)
} else {
arrays = append(arrays, t)
}
return t
}