mlx-bun 0.0.4

package/src/mlx/array.ts

@@ -0,0 +1,237 @@
+// MlxArray: ownership wrapper around an mlx_array handle.
+// Explicit .dispose() is the contract; a FinalizationRegistry backstop
+// frees leaked handles on GC (verified in spikes/phase0-memory.ts).
+
+import { JSCallback, ptr, toArrayBuffer } from "bun:ffi";
+import { C, Dtype, DTYPE_NAMES, type MlxHandle, optInt, outArray, takeMlxError } from "./ffi";
+import type { SafetensorsDtype } from "../safetensors";
+
+export const gpuStream: MlxHandle = C.mlx_default_gpu_stream_new();
+export const cpuStream: MlxHandle = C.mlx_default_cpu_stream_new();
+
+/** Shape buffer for FFI; ptr() rejects empty views, so 0-d arrays pass a
+ *  dummy buffer with dim=0 (the pointee is never read). */
+function shapeBuf(shape: number[]): Int32Array {
+  return shape.length === 0 ? new Int32Array(1) : new Int32Array(shape);
+}
+
+export const SAFETENSORS_TO_MLX: Record<SafetensorsDtype, Dtype> = {
+  BOOL: Dtype.bool, U8: Dtype.uint8, I8: Dtype.int8,
+  U16: Dtype.uint16, I16: Dtype.int16, U32: Dtype.uint32, I32: Dtype.int32,
+  U64: Dtype.uint64, I64: Dtype.int64,
+  F16: Dtype.float16, BF16: Dtype.bfloat16, F32: Dtype.float32, F64: Dtype.float64,
+};
+
+// --- zero-copy pinning ---------------------------------------------------
+// Buffers handed to mlx via new_data_managed_payload must outlive the mlx
+// array. We pin the JS view in a map keyed by a payload id; the shared dtor
+// callback unpins when mlx drops its last reference. threadsafe: mlx may
+// release buffers from its eval thread.
+
+const pinned = new Map<number, Uint8Array>();
+let nextPinId = 1;
+
+const unpinCallback = new JSCallback(
+  (payloadId: number) => { pinned.delete(payloadId); },
+  { args: ["u64"], returns: "void", threadsafe: true },
+);
+
+export function pinnedBufferCount(): number {
+  return pinned.size;
+}
+
+// --- finalization backstop ----------------------------------------------
+
+const registry = new FinalizationRegistry((handle: MlxHandle) => {
+  C.mlx_array_free(handle);
+});
+
+export class MlxArray {
+  #handle: MlxHandle;
+  #disposed = false;
+  #token = {};
+
+  constructor(handle: MlxHandle) {
+    this.#handle = handle;
+    registry.register(this, handle, this.#token);
+  }
+
+  /** Zero-copy: wrap an existing buffer (e.g. an mmap view). The view is
+   *  pinned until mlx releases it. */
+  static fromView(view: Uint8Array, shape: number[], dtype: Dtype): MlxArray {
+    const sb = shapeBuf(shape);
+    const id = nextPinId++;
+    pinned.set(id, view);
+    const handle = C.mlx_array_new_data_managed_payload(
+      ptr(view), ptr(sb), shape.length, dtype, id, unpinCallback.ptr,
+    );
+    return new MlxArray(handle);
+  }
+
+  /** Zero-copy: wrap a raw pointer (e.g. into an mmap'd weight shard).
+   *  The caller guarantees the memory outlives the array — weight mmaps
+   *  live for the process, so the dtor is a no-op unpin of id 0. */
+  static fromPointer(dataPtr: number, shape: number[], dtype: Dtype): MlxArray {
+    const sb = shapeBuf(shape);
+    const handle = C.mlx_array_new_data_managed_payload(
+      dataPtr, ptr(sb), shape.length, dtype, 0, unpinCallback.ptr,
+    );
+    return new MlxArray(handle);
+  }
+
+  /** Copying constructor for small host data. */
+  static fromFloat32(data: Float32Array, shape: number[]): MlxArray {
+    const sb = shapeBuf(shape);
+    return new MlxArray(
+      C.mlx_array_new_data(ptr(data), ptr(sb), shape.length, Dtype.float32),
+    );
+  }
+
+  /** Copy raw bytes into a fresh mlx-managed (page-aligned) LEAF array of the
+   *  given dtype/shape. Unlike fromView (which no-copy-wraps the host buffer and
+   *  is unsafe to feed to GPU ops when unaligned), mlx_array_new_data COPIES, so
+   *  the source bytes need not outlive the array and GPU ops are safe. Used to
+   *  detach a value from its computation graph: `MlxArray.fromBytesCopy(
+   *  x.rawBytes(), x.shape, x.dtype)` yields a graph-free leaf with x's exact
+   *  bytes (segmented-backward boundaries — src/train/segmented.ts). */
+  static fromBytesCopy(bytes: Uint8Array, shape: number[], dtype: Dtype): MlxArray {
+    const sb = shapeBuf(shape);
+    return new MlxArray(
+      C.mlx_array_new_data(ptr(bytes), ptr(sb), shape.length, dtype),
+    );
+  }
+
+  static fromInt32(data: Int32Array, shape: number[]): MlxArray {
+    const sb = shapeBuf(shape);
+    return new MlxArray(
+      C.mlx_array_new_data(ptr(data), ptr(sb), shape.length, Dtype.int32),
+    );
+  }
+
+  get handle(): MlxHandle {
+    if (this.#disposed) throw new Error("MlxArray used after dispose");
+    return this.#handle;
+  }
+
+  get ndim(): number {
+    return Number(C.mlx_array_ndim(this.handle));
+  }
+
+  // NOTE on the DFG stale-read bug (repro/bun-ffi-f64/ISSUE.md): the
+  // toArrayBuffer readbacks below (shape/rawBytes/toFloat32) are safe.
+  // The hazard is reading a *pre-existing* typed array after an FFI call
+  // wrote through its pointer — the JIT forwards a stale value from before
+  // the call. Here the view is constructed *after* the call, from the
+  // pointer the call returned: the loads are data-dependent on the call
+  // result and there is no prior JS load/store to forward from.
+
+  get shape(): number[] {
+    const n = this.ndim;
+    if (n === 0) return [];
+    const shapePtr = C.mlx_array_shape(this.handle);
+    return [...new Int32Array(toArrayBuffer(shapePtr!, 0, n * 4))];
+  }
+
+  get dtype(): Dtype {
+    return C.mlx_array_dtype(this.handle) as Dtype;
+  }
+
+  get dtypeName(): string {
+    return DTYPE_NAMES[this.dtype] ?? `unknown(${this.dtype})`;
+  }
+
+  get size(): number {
+    return Number(C.mlx_array_size(this.handle));
+  }
+
+  get nbytes(): number {
+    return Number(C.mlx_array_nbytes(this.handle));
+  }
+
+  eval(): this {
+    if (C.mlx_array_eval(this.handle) !== 0) throw new Error(`mlx_array_eval failed: ${takeMlxError() ?? ""}`);
+    return this;
+  }
+
+  /** Cast to another dtype (lazy). */
+  astype(dtype: Dtype, stream: MlxHandle = gpuStream): MlxArray {
+    return new MlxArray(
+      outArray("astype", (slot) => C.mlx_astype(slot, this.handle, dtype, stream)),
+    );
+  }
+
+  /** Raw bytes of the evaluated array (copy). bf16/f16/f32 only. */
+  rawBytes(): Uint8Array {
+    this.eval();
+    const dt = this.dtype;
+    const p =
+      dt === Dtype.float32 ? C.mlx_array_data_float32(this.handle)
+      : dt === Dtype.float16 ? C.mlx_array_data_float16(this.handle)
+      : dt === Dtype.bfloat16 ? C.mlx_array_data_bfloat16(this.handle)
+      : dt === Dtype.uint32 ? C.mlx_array_data_uint32(this.handle)
+      : null;
+    if (p === null) throw new Error(`rawBytes: unsupported dtype ${this.dtypeName}`);
+    return new Uint8Array(toArrayBuffer(p!, 0, this.nbytes)).slice();
+  }
+
+  /** Read back as float32 (casts on GPU if needed, then copies out). */
+  toFloat32(): Float32Array {
+    const src = this.dtype === Dtype.float32 ? this : this.astype(Dtype.float32);
+    src.eval();
+    const dataPtr = C.mlx_array_data_float32(src.handle);
+    const out = new Float32Array(toArrayBuffer(dataPtr!, 0, src.size * 4)).slice();
+    if (src !== this) src.dispose();
+    return out;
+  }
+
+  /** Basic slice: start/stop per dimension (stride 1). */
+  slice(start: number[], stop: number[], stream: MlxHandle = gpuStream): MlxArray {
+    const s = new Int32Array(start);
+    const e = new Int32Array(stop);
+    const strides = new Int32Array(start.map(() => 1));
+    return new MlxArray(
+      outArray("slice", (slot) =>
+        C.mlx_slice(
+          slot, this.handle,
+          ptr(s), BigInt(s.length), ptr(e), BigInt(e.length),
+          ptr(strides), BigInt(strides.length), stream,
+        ),
+      ),
+    );
+  }
+
+  add(other: MlxArray, stream: MlxHandle = gpuStream): MlxArray {
+    return new MlxArray(
+      outArray("add", (slot) => C.mlx_add(slot, this.handle, other.handle, stream)),
+    );
+  }
+
+  /** Dequantize (w, scales, biases) → full precision. */
+  static dequantize(
+    w: MlxArray, scales: MlxArray, biases: MlxArray | null,
+    groupSize: number, bits: number, mode = "affine",
+    stream: MlxHandle = gpuStream,
+  ): MlxArray {
+    const modeBuf = Buffer.from(mode + "\0", "utf8");
+    return new MlxArray(
+      outArray("dequantize", (slot) =>
+        C.mlx_dequantize(
+          slot, w.handle, scales.handle, biases?.handle ?? 0n,
+          optInt(groupSize), optInt(bits), ptr(modeBuf),
+          0n /* global_scale */, optInt(null) /* out dtype */, stream,
+        ),
+      ),
+    );
+  }
+
+  dispose(): void {
+    if (this.#disposed) return;
+    this.#disposed = true;
+    registry.unregister(this.#token);
+    C.mlx_array_free(this.#handle);
+  }
+
+  [Symbol.dispose](): void {
+    this.dispose();
+  }
+}

package/src/mlx/autograd.ts

@@ -0,0 +1,267 @@
+// Native autograd via mlx-c value_and_grad over a JS loss closure.
+//
+// Encapsulates the exact FFI lifecycle proven in 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";
+
+/** 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: JSCallback;
+  readonly #closure: bigint;
+  readonly #vag: bigint;
+  readonly #argIdx: Int32Array;
+  #closureError: string | null = null;
+  #disposed = false;
+
+  constructor(loss: (primals: MlxArray[]) => MlxArray, argIdx: number[]) {
+    if (argIdx.length === 0) throw new Error("ValueAndGrad: argIdx must be non-empty");
+    this.#argIdx = new Int32Array(argIdx);
+
+    this.#cb = new JSCallback(
+      (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 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" },
+    );
+
+    this.#closure = C.mlx_closure_new_func_payload(this.#cb.ptr as never, null, null);
+
+    const vagSlot = new BigUint64Array([C.mlx_closure_value_and_grad_new()]);
+    if (C.mlx_value_and_grad(ptr(vagSlot), this.#closure, ptr(this.#argIdx), BigInt(this.#argIdx.length)) !== 0) {
+      // best-effort cleanup before throwing
+      C.mlx_closure_free(this.#closure);
+      this.#cb.close();
+      throw new Error("mlx_value_and_grad failed");
+    }
+    this.#vag = read.u64(ptr(vagSlot), 0);
+  }
+
+  /** 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;
+    C.mlx_closure_value_and_grad_free(this.#vag);
+    C.mlx_closure_free(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

@@ -0,0 +1,116 @@
+// 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";
+
+/** 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: JSCallback;
+  readonly #closure: bigint;
+  readonly #ckpt: bigint;
+  #closureError: string | null = null;
+  #disposed = false;
+
+  constructor(fn: (inputs: MlxArray[]) => MlxArray[]) {
+    this.#cb = new JSCallback(
+      (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" },
+    );
+
+    this.#closure = C.mlx_closure_new_func_payload(this.#cb.ptr as never, null, null);
+
+    const slot = new BigUint64Array([C.mlx_closure_new()]);
+    if (C.mlx_checkpoint(ptr(slot), this.#closure) !== 0) {
+      C.mlx_closure_free(this.#closure);
+      this.#cb.close();
+      throw new Error(`mlx_checkpoint failed: ${takeMlxError() ?? ""}`);
+    }
+    this.#ckpt = read.u64(ptr(slot), 0);
+  }
+
+  /** 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;
+    C.mlx_closure_free(this.#ckpt);
+    C.mlx_closure_free(this.#closure);
+    this.#cb.close();
+  }
+
+  [Symbol.dispose](): void {
+    this.dispose();
+  }
+}