mirror of
https://github.com/ollama/ollama.git
synced 2026-04-20 15:55:46 +02:00
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:
@@ -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
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user