mlx-bun 0.5.0 → 1.0.0

package/src/mlx/autograd.ts

@@ -1,304 +0,0 @@
-// Native autograd via mlx-c value_and_grad over a JS loss closure.
-//
-// Encapsulates the exact FFI lifecycle proven in lab/spikes/phase-train-vag.ts:
-//   JS loss closure (JSCallback) → mlx_closure_new_func_payload →
-//   mlx_value_and_grad → mlx_closure_value_and_grad_apply → read grads.
-//
-// House rules honored:
-// - Out-params read via bun:ffi `read.u64` (never slot[0] — DFG stale-read).
-// - The JS closure NEVER throws across the FFI boundary (try/catch → return 1);
-//   the captured error is re-thrown on the JS side after apply() fails.
-// - All native handles (closure, value_and_grad closure, JSCallback) are held
-//   alive on the instance and freed in dispose().
-
-import { JSCallback, ptr, read } from "bun:ffi";
-import { C, takeMlxError } from "./ffi";
-import { MlxArray } from "./array";
-
-type CallbackHandle = Pick<JSCallback, "ptr" | "close">;
-type NativeCall<F extends (...args: any[]) => any> =
-  (...args: Parameters<F>) => ReturnType<F>;
-
-/** Constructor-only dependency seam for deterministic native failure tests. */
-export interface ValueAndGradConstructorHooks {
-  callback?: (
-    callback: ConstructorParameters<typeof JSCallback>[0],
-    definition: ConstructorParameters<typeof JSCallback>[1],
-  ) => CallbackHandle;
-  closureNew?: NativeCall<typeof C.mlx_closure_new_func_payload>;
-  transformedNew?: NativeCall<typeof C.mlx_closure_value_and_grad_new>;
-  transform?: NativeCall<typeof C.mlx_value_and_grad>;
-  transformedFree?: NativeCall<typeof C.mlx_closure_value_and_grad_free>;
-  closureFree?: NativeCall<typeof C.mlx_closure_free>;
-}
-
-/** A reusable value-and-grad over a JS-built loss graph.
- *
- *  `loss(primals)` receives the input arrays in argument order and must
- *  return a single SCALAR MlxArray (the loss). The graph it builds is
- *  differentiated w.r.t. the inputs named in `argIdx`. `apply(primals)`
- *  returns the loss value and exactly `argIdx.length` gradients (one per
- *  differentiated input, in `argIdx` order).
- *
- *  The closure must not throw; if the user `loss` throws, the error is
- *  captured and surfaced from `apply()`. The MlxArrays handed to `loss`
- *  are owned by mlx — do NOT dispose them; the returned loss IS disposed
- *  of by mlx's graph. */
-export class ValueAndGrad {
-  readonly #cb: CallbackHandle;
-  readonly #closure: bigint;
-  readonly #vag: bigint;
-  readonly #argIdx: Int32Array;
-  readonly #freeClosure: NativeCall<typeof C.mlx_closure_free>;
-  readonly #freeVag: NativeCall<typeof C.mlx_closure_value_and_grad_free>;
-  #closureError: string | null = null;
-  #disposed = false;
-
-  constructor(
-    loss: (primals: MlxArray[]) => MlxArray,
-    argIdx: number[],
-    hooks: ValueAndGradConstructorHooks = {},
-  ) {
-    if (argIdx.length === 0) throw new Error("ValueAndGrad: argIdx must be non-empty");
-    this.#argIdx = new Int32Array(argIdx);
-    const makeCallback = hooks.callback ?? ((callback, definition) => new JSCallback(callback, definition));
-    const closureNew = hooks.closureNew ?? C.mlx_closure_new_func_payload;
-    const transformedNew = hooks.transformedNew ?? C.mlx_closure_value_and_grad_new;
-    const transform = hooks.transform ?? C.mlx_value_and_grad;
-    this.#freeClosure = hooks.closureFree ?? C.mlx_closure_free;
-    this.#freeVag = hooks.transformedFree ?? C.mlx_closure_value_and_grad_free;
-
-    const cb = makeCallback(
-      (outPtr: number, inVec: bigint, _payload: number): number => {
-        try {
-          // Read the primals out of the input vector_array. mlx_vector_array_get
-          // returns a fresh (ref-counted) handle, so disposing the wrapper
-          // afterward is correct — it drops our reference, not mlx's graph
-          // reference (proven in lab/spikes/phase-train-vag.ts).
-          const n = Number(C.mlx_vector_array_size(inVec));
-          const primals: MlxArray[] = [];
-          for (let i = 0; i < n; i++) {
-            const slot = new BigUint64Array([C.mlx_array_new()]);
-            const sp = ptr(slot);
-            if (C.mlx_vector_array_get(sp, inVec, BigInt(i)) !== 0)
-              throw new Error(`vector_array_get(${i}) failed`);
-            primals.push(new MlxArray(read.u64(sp, 0)));
-          }
-          const out = loss(primals);
-          // Populate *out with the single scalar loss. set_data copies the
-          // handle into the output vector (mlx retains its own reference), so
-          // disposing the loss wrapper afterward drops only our reference.
-          const handles = new BigUint64Array([out.handle]);
-          if (C.mlx_vector_array_set_data(outPtr as never, ptr(handles), 1n) !== 0)
-            throw new Error("vector_array_set_data failed");
-          for (const p of primals) p.dispose();
-          out.dispose();
-          return 0;
-        } catch (e) {
-          this.#closureError = e instanceof Error ? e.message : String(e);
-          return 1;
-        }
-      },
-      { args: ["ptr", "u64", "ptr"], returns: "i32" },
-    );
-
-    let closure: bigint | null = null;
-    let vagSlot: BigUint64Array | null = null;
-    try {
-      closure = closureNew(cb.ptr as never, null, null);
-      vagSlot = new BigUint64Array([transformedNew()]);
-      if (transform(ptr(vagSlot), closure, ptr(this.#argIdx), BigInt(this.#argIdx.length)) !== 0)
-        throw new Error("mlx_value_and_grad failed");
-      this.#cb = cb;
-      this.#closure = closure;
-      this.#vag = read.u64(ptr(vagSlot), 0);
-    } catch (e) {
-      // mlx-c out-params are initialized live handles. On failure the slot
-      // still owns whichever handle it contains and must be freed explicitly.
-      if (vagSlot) this.#freeVag(read.u64(ptr(vagSlot), 0));
-      if (closure !== null) this.#freeClosure(closure);
-      cb.close();
-      throw e;
-    }
-  }
-
-  /** Run the loss + backward on `primals` (in argument order). Returns the
-   *  scalar loss and one gradient per differentiated input. Caller owns all
-   *  returned arrays (dispose them). */
-  apply(primals: MlxArray[]): { value: MlxArray; grads: MlxArray[] } {
-    if (this.#disposed) throw new Error("ValueAndGrad used after dispose");
-    this.#closureError = null;
-
-    const handles = new BigUint64Array(primals.map((p) => p.handle));
-    const inVec = C.mlx_vector_array_new_data(ptr(handles), BigInt(primals.length));
-
-    const valueSlot = new BigUint64Array([C.mlx_vector_array_new()]);
-    const gradsSlot = new BigUint64Array([C.mlx_vector_array_new()]);
-    const st = C.mlx_closure_value_and_grad_apply(ptr(valueSlot), ptr(gradsSlot), this.#vag, inVec);
-    const valueVec = read.u64(ptr(valueSlot), 0);
-    const gradsVec = read.u64(ptr(gradsSlot), 0);
-    if (st !== 0) {
-      C.mlx_vector_array_free(valueVec);
-      C.mlx_vector_array_free(gradsVec);
-      C.mlx_vector_array_free(inVec);
-      // Surface whichever error fired: a JS throw inside the loss closure
-      // (#closureError) or MLX's own C-level error (takeMlxError) — e.g. a
-      // missing vjp for an op in the forward graph. Discarding the latter
-      // turned an "unsupported gradient" into an opaque failure.
-      const mlxErr = takeMlxError();
-      const detail = this.#closureError ?? mlxErr ?? "";
-      throw new Error(`value_and_grad_apply failed${detail ? `: ${detail}` : ""}`);
-    }
-
-    // value vector has exactly one entry (the scalar loss).
-    const value = vecGet(valueVec, 0);
-
-    const nGrads = Number(C.mlx_vector_array_size(gradsVec));
-    if (nGrads !== this.#argIdx.length) {
-      value.dispose();
-      C.mlx_vector_array_free(valueVec);
-      C.mlx_vector_array_free(gradsVec);
-      C.mlx_vector_array_free(inVec);
-      throw new Error(`expected ${this.#argIdx.length} grads, got ${nGrads}`);
-    }
-    const grads: MlxArray[] = [];
-    for (let i = 0; i < nGrads; i++) grads.push(vecGet(gradsVec, i));
-
-    C.mlx_vector_array_free(valueVec);
-    C.mlx_vector_array_free(gradsVec);
-    C.mlx_vector_array_free(inVec);
-    return { value, grads };
-  }
-
-  dispose(): void {
-    if (this.#disposed) return;
-    this.#disposed = true;
-    this.#freeVag(this.#vag);
-    this.#freeClosure(this.#closure);
-    this.#cb.close();
-  }
-
-  [Symbol.dispose](): void {
-    this.dispose();
-  }
-}
-
-/** Pull array `i` out of a vector_array into an owned MlxArray wrapper. */
-function vecGet(vec: bigint, i: number): MlxArray {
-  const slot = new BigUint64Array([C.mlx_array_new()]);
-  const sp = ptr(slot);
-  if (C.mlx_vector_array_get(sp, vec, BigInt(i)) !== 0)
-    throw new Error(`vector_array_get(${i}) failed`);
-  return new MlxArray(read.u64(sp, 0));
-}
-
-/** Reverse-mode AD via mlx_vjp: given a function `fn(primals) -> outputs` and a
- *  cotangent per output, returns the cotangent w.r.t. EVERY primal. This is the
- *  natural primitive for segmented backward (segment forward + dh -> [dh_in,
- *  ...dLoRA]) — no surrogate scalar and no value_and_grad closure. Unlike
- *  ValueAndGrad it returns a vjp for ALL primals (not a selected argnum subset).
- *
- *  `fn` must build its output graph from the input arrays in argument order and
- *  return them (most often a single array). It must not throw; a throw is
- *  captured and surfaced from apply(). The MlxArrays handed to `fn` are owned by
- *  mlx — do NOT dispose them. Lifecycle/house-rules mirror ValueAndGrad. */
-export class Vjp {
-  readonly #cb: JSCallback;
-  readonly #closure: bigint;
-  readonly #nOut: number;
-  #closureError: string | null = null;
-  #disposed = false;
-
-  /** @param nOut number of arrays `fn` returns (length of the cotangent list). */
-  constructor(fn: (primals: MlxArray[]) => MlxArray[], nOut = 1) {
-    this.#nOut = nOut;
-    this.#cb = new JSCallback(
-      (outPtr: number, inVec: bigint, _payload: number): number => {
-        try {
-          const n = Number(C.mlx_vector_array_size(inVec));
-          const primals: MlxArray[] = [];
-          for (let i = 0; i < n; i++) {
-            const slot = new BigUint64Array([C.mlx_array_new()]);
-            const sp = ptr(slot);
-            if (C.mlx_vector_array_get(sp, inVec, BigInt(i)) !== 0)
-              throw new Error(`vector_array_get(${i}) failed`);
-            primals.push(new MlxArray(read.u64(sp, 0)));
-          }
-          const outs = fn(primals);
-          const handles = new BigUint64Array(outs.map((o) => o.handle));
-          if (C.mlx_vector_array_set_data(outPtr as never, ptr(handles), BigInt(outs.length)) !== 0)
-            throw new Error("vector_array_set_data failed");
-          for (const p of primals) p.dispose();
-          for (const o of outs) o.dispose();
-          return 0;
-        } catch (e) {
-          this.#closureError = e instanceof Error ? e.message : String(e);
-          return 1;
-        }
-      },
-      { args: ["ptr", "u64", "ptr"], returns: "i32" },
-    );
-    this.#closure = C.mlx_closure_new_func_payload(this.#cb.ptr as never, null, null);
-  }
-
-  /** Run the vjp on `primals` with `cotangents` (one per output of `fn`).
-   *  Returns the outputs (fn(primals)) and one cotangent per primal (in primal
-   *  order). Caller owns all returned arrays. */
-  apply(primals: MlxArray[], cotangents: MlxArray[]): { outputs: MlxArray[]; vjps: MlxArray[] } {
-    if (this.#disposed) throw new Error("Vjp used after dispose");
-    if (cotangents.length !== this.#nOut)
-      throw new Error(`Vjp: expected ${this.#nOut} cotangents, got ${cotangents.length}`);
-    this.#closureError = null;
-
-    const pHandles = new BigUint64Array(primals.map((p) => p.handle));
-    const cHandles = new BigUint64Array(cotangents.map((c) => c.handle));
-    const inVec = C.mlx_vector_array_new_data(ptr(pHandles), BigInt(primals.length));
-    const cotVec = C.mlx_vector_array_new_data(ptr(cHandles), BigInt(cotangents.length));
-
-    const outSlot = new BigUint64Array([C.mlx_vector_array_new()]);
-    const vjpSlot = new BigUint64Array([C.mlx_vector_array_new()]);
-    const st = C.mlx_vjp(ptr(outSlot), ptr(vjpSlot), this.#closure, inVec, cotVec);
-    const outVec = read.u64(ptr(outSlot), 0);
-    const vjpVec = read.u64(ptr(vjpSlot), 0);
-    if (st !== 0) {
-      C.mlx_vector_array_free(outVec);
-      C.mlx_vector_array_free(vjpVec);
-      C.mlx_vector_array_free(inVec);
-      C.mlx_vector_array_free(cotVec);
-      const mlxErr = takeMlxError();
-      const detail = this.#closureError ?? mlxErr ?? "";
-      throw new Error(`mlx_vjp failed${detail ? `: ${detail}` : ""}`);
-    }
-
-    const nVjp = Number(C.mlx_vector_array_size(vjpVec));
-    if (nVjp !== primals.length) {
-      C.mlx_vector_array_free(outVec);
-      C.mlx_vector_array_free(vjpVec);
-      C.mlx_vector_array_free(inVec);
-      C.mlx_vector_array_free(cotVec);
-      throw new Error(`Vjp: expected ${primals.length} vjps, got ${nVjp}`);
-    }
-    const nOut = Number(C.mlx_vector_array_size(outVec));
-    const outputs: MlxArray[] = [];
-    for (let i = 0; i < nOut; i++) outputs.push(vecGet(outVec, i));
-    const vjps: MlxArray[] = [];
-    for (let i = 0; i < nVjp; i++) vjps.push(vecGet(vjpVec, i));
-
-    C.mlx_vector_array_free(outVec);
-    C.mlx_vector_array_free(vjpVec);
-    C.mlx_vector_array_free(inVec);
-    C.mlx_vector_array_free(cotVec);
-    return { outputs, vjps };
-  }
-
-  dispose(): void {
-    if (this.#disposed) return;
-    this.#disposed = true;
-    C.mlx_closure_free(this.#closure);
-    this.#cb.close();
-  }
-
-  [Symbol.dispose](): void {
-    this.dispose();
-  }
-}

package/src/mlx/checkpoint.ts

@@ -1,149 +0,0 @@
-// Gradient checkpointing (rematerialization) via mlx-c mlx_checkpoint.
-//
-// Wraps a JS function (MlxArray[]) -> MlxArray[] into a closure whose forward
-// activations are dropped and RECOMPUTED during the enclosing value_and_grad
-// backward. Used to bound training memory on long sequences: instead of
-// retaining every transformer layer's activations, each checkpointed unit is
-// recomputed one at a time in the backward pass.
-//
-// Same FFI lifecycle as ValueAndGrad (src/mlx/autograd.ts): a JSCallback over
-// the user fn → mlx_closure_new_func_payload → mlx_checkpoint → applied with
-// mlx_closure_apply. The closure must not throw across the FFI boundary; a JS
-// throw is captured and re-surfaced from apply().
-
-import { JSCallback, ptr, read } from "bun:ffi";
-import { C, takeMlxError } from "./ffi";
-import { MlxArray } from "./array";
-
-type CallbackHandle = Pick<JSCallback, "ptr" | "close">;
-type NativeCall<F extends (...args: any[]) => any> =
-  (...args: Parameters<F>) => ReturnType<F>;
-
-/** Constructor-only dependency seam for deterministic native failure tests. */
-export interface CheckpointConstructorHooks {
-  callback?: (
-    callback: ConstructorParameters<typeof JSCallback>[0],
-    definition: ConstructorParameters<typeof JSCallback>[1],
-  ) => CallbackHandle;
-  closureNew?: NativeCall<typeof C.mlx_closure_new_func_payload>;
-  transformedNew?: NativeCall<typeof C.mlx_closure_new>;
-  transform?: NativeCall<typeof C.mlx_checkpoint>;
-  closureFree?: NativeCall<typeof C.mlx_closure_free>;
-}
-
-/** A reusable checkpointed wrapper over a JS-built subgraph.
- *
- *  `fn(inputs)` receives the input arrays (in order) and returns one or more
- *  output arrays. Applying the checkpoint builds the same forward graph, but
- *  marks it for recomputation: the enclosing backward recomputes `fn` rather
- *  than reading retained activations. Gradients flow to the explicit inputs
- *  AND to any outer-trace leaves captured by `fn` (e.g. swapped-in LoRA
- *  primals), exactly as if `fn` had been inlined. */
-export class Checkpoint {
-  readonly #cb: CallbackHandle;
-  readonly #closure: bigint;
-  readonly #ckpt: bigint;
-  readonly #freeClosure: NativeCall<typeof C.mlx_closure_free>;
-  #closureError: string | null = null;
-  #disposed = false;
-
-  constructor(
-    fn: (inputs: MlxArray[]) => MlxArray[],
-    hooks: CheckpointConstructorHooks = {},
-  ) {
-    const makeCallback = hooks.callback ?? ((callback, definition) => new JSCallback(callback, definition));
-    const closureNew = hooks.closureNew ?? C.mlx_closure_new_func_payload;
-    const transformedNew = hooks.transformedNew ?? C.mlx_closure_new;
-    const transform = hooks.transform ?? C.mlx_checkpoint;
-    this.#freeClosure = hooks.closureFree ?? C.mlx_closure_free;
-
-    const cb = makeCallback(
-      (outPtr: number, inVec: bigint, _payload: number): number => {
-        try {
-          const n = Number(C.mlx_vector_array_size(inVec));
-          const inputs: MlxArray[] = [];
-          for (let i = 0; i < n; i++) {
-            const slot = new BigUint64Array([C.mlx_array_new()]);
-            const sp = ptr(slot);
-            if (C.mlx_vector_array_get(sp, inVec, BigInt(i)) !== 0)
-              throw new Error(`vector_array_get(${i}) failed`);
-            inputs.push(new MlxArray(read.u64(sp, 0)));
-          }
-          const outs = fn(inputs);
-          const handles = new BigUint64Array(outs.map((o) => o.handle));
-          if (C.mlx_vector_array_set_data(outPtr as never, ptr(handles), BigInt(outs.length)) !== 0)
-            throw new Error("vector_array_set_data failed");
-          for (const p of inputs) p.dispose();
-          for (const o of outs) o.dispose();
-          return 0;
-        } catch (e) {
-          this.#closureError = e instanceof Error ? e.message : String(e);
-          return 1;
-        }
-      },
-      { args: ["ptr", "u64", "ptr"], returns: "i32" },
-    );
-
-    let closure: bigint | null = null;
-    let slot: BigUint64Array | null = null;
-    try {
-      closure = closureNew(cb.ptr as never, null, null);
-      slot = new BigUint64Array([transformedNew()]);
-      if (transform(ptr(slot), closure) !== 0)
-        throw new Error(`mlx_checkpoint failed: ${takeMlxError() ?? ""}`);
-      this.#cb = cb;
-      this.#closure = closure;
-      this.#ckpt = read.u64(ptr(slot), 0);
-    } catch (e) {
-      if (slot) this.#freeClosure(read.u64(ptr(slot), 0));
-      if (closure !== null) this.#freeClosure(closure);
-      cb.close();
-      throw e;
-    }
-  }
-
-  /** Apply the checkpointed fn. Returns its outputs as owned MlxArrays. */
-  apply(inputs: MlxArray[]): MlxArray[] {
-    if (this.#disposed) throw new Error("Checkpoint used after dispose");
-    this.#closureError = null;
-
-    const handles = new BigUint64Array(inputs.map((p) => p.handle));
-    const inVec = C.mlx_vector_array_new_data(ptr(handles), BigInt(inputs.length));
-    const outSlot = new BigUint64Array([C.mlx_vector_array_new()]);
-    const st = C.mlx_closure_apply(ptr(outSlot), this.#ckpt, inVec);
-    const outVec = read.u64(ptr(outSlot), 0);
-    if (st !== 0) {
-      C.mlx_vector_array_free(outVec);
-      C.mlx_vector_array_free(inVec);
-      const detail = this.#closureError ?? takeMlxError() ?? "";
-      throw new Error(`checkpoint apply failed${detail ? `: ${detail}` : ""}`);
-    }
-    const nOut = Number(C.mlx_vector_array_size(outVec));
-    const outs: MlxArray[] = [];
-    for (let i = 0; i < nOut; i++) {
-      const slot = new BigUint64Array([C.mlx_array_new()]);
-      const sp = ptr(slot);
-      if (C.mlx_vector_array_get(sp, outVec, BigInt(i)) !== 0) {
-        C.mlx_vector_array_free(outVec);
-        C.mlx_vector_array_free(inVec);
-        throw new Error(`checkpoint output get(${i}) failed`);
-      }
-      outs.push(new MlxArray(read.u64(sp, 0)));
-    }
-    C.mlx_vector_array_free(outVec);
-    C.mlx_vector_array_free(inVec);
-    return outs;
-  }
-
-  dispose(): void {
-    if (this.#disposed) return;
-    this.#disposed = true;
-    this.#freeClosure(this.#ckpt);
-    this.#freeClosure(this.#closure);
-    this.#cb.close();
-  }
-
-  [Symbol.dispose](): void {
-    this.dispose();
-  }
-}

package/src/mlx/compile.ts

@@ -1,168 +0,0 @@
-// mx.compile from Bun: wrap a JS graph-builder in an mlx_closure, hand it
-// to mlx_compile, and replay the compiled graph in C++ thereafter.
-//
-// The trace function runs EXACTLY ONCE per (ndim, dtype) input signature
-// (shapeless=true): mlx invokes our JSCallback with tracer arrays during
-// the first mlx_closure_apply, we build the graph with the ordinary ops
-// layer, and every later apply replays it natively — zero per-op FFI.
-//
-// Contract for trace functions:
-// - receive positional MlxArray inputs, return positional MlxArray outputs
-// - must NOT dispose the inputs (the trampoline owns them) and must NOT
-//   eval anything (tracer arrays have no data — eval inside a trace is an
-//   mlx error by design)
-// - arrays created inside the trace (weights, scalarLike constants) are
-//   captured as graph constants — correct for weights, which never change
-// - with shapeless=true the graph must stay valid as dim SIZES change
-//   (ndim/dtype changes retrace automatically): no reshape/slice constants
-//   derived from a varying dimension
-
-import { JSCallback, ptr, read } from "bun:ffi";
-import { C, type MlxHandle, takeMlxError } from "./ffi";
-import { MlxArray } from "./array";
-
-export type TraceFn = (inputs: MlxArray[]) => MlxArray[];
-
-const traceRegistry = new Map<number, { fn: TraceFn; count: number; shapeless: boolean }>();
-let nextTraceId = 1;
-let shapelessTracing = false;
-
-/** Fixed-shape custom kernels must fall back while a shapeless graph traces. */
-export function isShapelessTracing(): boolean {
-  return shapelessTracing;
-}
-/** Error thrown by a trace fn, carried across the C boundary (the
- *  trampoline must not throw into mlx; it returns non-zero instead). */
-let traceError: unknown = null;
-
-/** Count of trace invocations, exposed for tests (a working shapeless
- *  compile traces once and replays thereafter). */
-export let traceCalls = 0;
-
-// int (*fun)(mlx_vector_array* res, const mlx_vector_array args, void* payload)
-// mlx_vector_array is the usual one-pointer struct, by value → u64.
-// Tracing happens synchronously inside mlx_closure_apply on the calling
-// thread, so a non-threadsafe JSCallback is safe (same argument as the
-// error handler in ffi.ts).
-const trampoline = new JSCallback(
-  (resPtr: number, argsVec: bigint, payload: number): number => {
-    const previousShapeless = shapelessTracing;
-    try {
-      const entry = traceRegistry.get(payload);
-      if (!entry) throw new Error(`compile trampoline: unknown trace id ${payload}`);
-      shapelessTracing ||= entry.shapeless;
-      const fn = entry.fn;
-      entry.count++;
-      traceCalls++;
-      const n = Number(C.mlx_vector_array_size(argsVec));
-      const inputs: MlxArray[] = [];
-      const slot = new BigUint64Array(1);
-      const slotPtr = ptr(slot);
-      for (let i = 0; i < n; i++) {
-        slot[0] = C.mlx_array_new();
-        if (C.mlx_vector_array_get(slotPtr, argsVec, BigInt(i)) !== 0)
-          throw new Error(`vector_array_get(${i}) failed: ${takeMlxError() ?? ""}`);
-        inputs.push(new MlxArray(read.u64(slotPtr, 0)));
-      }
-      const outputs = fn(inputs);
-      const handles = new BigUint64Array(outputs.map((o) => o.handle));
-      if (C.mlx_vector_array_set_data(resPtr, ptr(handles), BigInt(outputs.length)) !== 0)
-        throw new Error(`vector_array_set_data failed: ${takeMlxError() ?? ""}`);
-      for (const o of outputs) o.dispose();
-      for (const a of inputs) a.dispose();
-      return 0;
-    } catch (e) {
-      traceError = e;
-      return 1;
-    } finally {
-      shapelessTracing = previousShapeless;
-    }
-  },
-  { args: ["ptr", "u64", "ptr"], returns: "i32" },
-);
-
-/** MLX_COMPILE_MODE_* (compile.h). NO_FUSE keeps the host-side replay win
- *  while disabling kernel fusion — the escape hatch if fusion ever breaks
- *  the bit-exact gate. */
-export function setCompileMode(mode: "disabled" | "no_simplify" | "no_fuse" | "enabled"): void {
-  const v = { disabled: 0, no_simplify: 1, no_fuse: 2, enabled: 3 }[mode];
-  if (C.mlx_set_compile_mode(v) !== 0)
-    throw new Error(`mlx_set_compile_mode failed: ${takeMlxError() ?? ""}`);
-}
-
-export class CompiledFunction {
-  #compiled: MlxHandle;
-  #traceId: number;
-  #disposed = false;
-
-  /** Times this function's trace ran. A shapeless decode closure must
-   *  trace exactly once — its input ndims/dtypes never change — so any
-   *  count > 1 is a silent retrace (per-step JS graph builds sneaking
-   *  back in) and gets flagged by compiled-decode. */
-  get traceCount(): number {
-    return traceRegistry.get(this.#traceId)?.count ?? 0;
-  }
-
-  constructor(fn: TraceFn, shapeless = true) {
-    this.#traceId = nextTraceId++;
-    traceRegistry.set(this.#traceId, { fn, count: 0, shapeless });
-    // payload is the trace id smuggled through the void* (never deref'd);
-    // dtor is null — nothing to free on the C side.
-    const raw = C.mlx_closure_new_func_payload(trampoline.ptr, this.#traceId as never, null);
-    const slot = new BigUint64Array([C.mlx_closure_new()]);
-    const slotPtr = ptr(slot);
-    const status = C.mlx_compile(slotPtr, raw, shapeless);
-    const compiled = read.u64(slotPtr, 0);
-    C.mlx_closure_free(raw);
-    if (status !== 0) {
-      C.mlx_closure_free(compiled);
-      traceRegistry.delete(this.#traceId);
-      throw new Error(`mlx_compile failed: ${takeMlxError() ?? ""}`);
-    }
-    this.#compiled = compiled;
-  }
-
-  /** Run the compiled graph. First call (per ndim/dtype signature) traces
-   *  through the JS fn; later calls replay natively. Returned arrays are
-   *  owned by the caller. */
-  apply(inputs: MlxArray[]): MlxArray[] {
-    if (this.#disposed) throw new Error("CompiledFunction used after dispose");
-    const handles = new BigUint64Array(inputs.length);
-    for (let i = 0; i < inputs.length; i++) handles[i] = inputs[i]!.handle;
-    const inVec = C.mlx_vector_array_new_data(ptr(handles), BigInt(inputs.length));
-    const slot = new BigUint64Array([C.mlx_vector_array_new()]);
-    const slotPtr = ptr(slot);
-    const status = C.mlx_closure_apply(slotPtr, this.#compiled, inVec);
-    const outVec = read.u64(slotPtr, 0);
-    C.mlx_vector_array_free(inVec);
-    if (status !== 0) {
-      C.mlx_vector_array_free(outVec);
-      const e = traceError;
-      traceError = null;
-      throw e instanceof Error
-        ? e
-        : new Error(`compiled apply failed: ${takeMlxError() ?? "(no mlx message)"}`);
-    }
-    const n = Number(C.mlx_vector_array_size(outVec));
-    const outputs: MlxArray[] = [];
-    const aSlot = new BigUint64Array(1);
-    const aPtr = ptr(aSlot);
-    for (let i = 0; i < n; i++) {
-      aSlot[0] = C.mlx_array_new();
-      if (C.mlx_vector_array_get(aPtr, outVec, BigInt(i)) !== 0) {
-        C.mlx_vector_array_free(outVec);
-        throw new Error(`vector_array_get(out ${i}) failed: ${takeMlxError() ?? ""}`);
-      }
-      outputs.push(new MlxArray(read.u64(aPtr, 0)));
-    }
-    C.mlx_vector_array_free(outVec);
-    return outputs;
-  }
-
-  dispose(): void {
-    if (this.#disposed) return;
-    this.#disposed = true;
-    C.mlx_closure_free(this.#compiled);
-    traceRegistry.delete(this.#traceId);
-  }
-}