mlx-bun 0.4.0 → 1.0.0

package/src/mlx/array.ts

@@ -1,388 +0,0 @@
-// MlxArray: ownership wrapper around an mlx_array handle.
-// Explicit .dispose() is the contract; a FinalizationRegistry backstop
-// frees leaked handles on GC (verified in lab/spikes/phase0-memory.ts).
-
-import { dlopen, 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. fromView pins the JS view in this map (a GC root, nothing more);
-// release is EXPLICIT and JS-thread-only via unpinHostBuffer() — never an
-// mlx dtor. See the hazard note below.
-
-const pinned = new Map<number, Uint8Array>();
-let nextPinId = 1;
-
-// HAZARD (root cause of the 2026-07-06 restart-restore hang): mlx releases
-// buffer Data wherever the LAST shared_ptr drops — often the Metal
-// COMPLETION thread (gpu::eval retains buffers until the command buffer
-// finishes, i.e. past dispose()). A bun:ffi JSCallback dtor invoked from
-// that thread deadlocks when the JS thread is inside a blocking FFI eval
-// (completion waits on JS, JS waits on completion) and SIGTRAPs mid-GC.
-// So NO array here ever gets a JS dtor: fromPointer AND fromView both hand
-// mlx the native no-op dtor below (payload 0), and host-buffer lifetime is
-// the CALLER's contract, released only from the JS thread (unpinHostBuffer)
-// or never (process-lifetime mmaps: expert-offload).
-
-// Native no-op dtor: payload is always 0, and libc free(NULL) is defined to
-// do nothing — a dtor mlx can safely call from ANY thread. The pointed-to
-// memory's lifetime is the CALLER's contract (weight mmaps live for the
-// process; restored KV is COPIED at load since 2026-07-07 — see kv-store.ts
-// loadKvCache — so no restore mapping needs pinning anymore). Bun >= 1.4
-// returns a numeric function address from .ptr; verified against dlsym on
-// 1.4.0 and 1.4.2. Keep the library open for every native destructor's lifetime.
-const libcFree = dlopen("/usr/lib/libSystem.B.dylib", {
-  free: { args: ["ptr"], returns: "void" },
-});
-// The pinned bun-types predates the runtime's numeric .ptr property.
-const freeFnAddr = Number(Reflect.get(libcFree.symbols.free, "ptr"));
-if (!Number.isSafeInteger(freeFnAddr) || freeFnAddr <= 0)
-  throw new Error("free.ptr is invalid — zero-copy arrays need a native no-op dtor");
-
-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 = {};
-  #pinId = 0; // fromView only: key into the pinned-view map (0 = no pin)
-
-  constructor(handle: MlxHandle) {
-    this.#handle = handle;
-    registry.register(this, handle, this.#token);
-  }
-
-  /** Zero-copy: wrap an existing JS buffer. The view is pinned (GC-rooted)
-   *  in a process-side map; mlx gets the NATIVE no-op dtor (free(NULL)) —
-   *  never a JS callback, so the last release may happen on the Metal
-   *  completion thread without deadlock (see hazard note above). Buffer
-   *  lifetime is therefore the CALLER's contract:
-   *   - mmap-backed views (toArrayBuffer aliases): the mapping owns the
-   *     memory; prefer fromPointer, which skips the pin entirely.
-   *   - JS-heap buffers: the pin keeps the buffer alive; call
-   *     unpinHostBuffer() only after a JS-side sync point proves mlx is
-   *     done with it (everything depending on the array evaluated AND the
-   *     array disposed), or never (process-lifetime, e.g. short scripts). */
-  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, 0, freeFnAddr,
-    );
-    const arr = new MlxArray(handle);
-    arr.#pinId = id;
-    return arr;
-  }
-
-  /** Release the GC root fromView placed on this array's host buffer.
-   *  JS-thread-only by construction (it's a plain method) — this is the
-   *  replacement for the deadlocking JSCallback unpin dtor. Idempotent;
-   *  no-op for non-fromView arrays. Caller contract: only call once mlx can
-   *  no longer touch the buffer (see fromView doc), or when the memory is
-   *  owned elsewhere (an mmap kept mapped past every use). */
-  unpinHostBuffer(): void {
-    if (this.#pinId !== 0) {
-      pinned.delete(this.#pinId);
-      this.#pinId = 0;
-    }
-  }
-
-  /** Zero-copy: wrap a raw pointer (e.g. into an mmap'd weight shard).
-   *  The caller guarantees the memory outlives EVERY mlx
-   *  reference to the array — including GPU command buffers that retain the
-   *  buffer past dispose() — i.e. mmaps backing these arrays stay mapped for
-   *  the process. The dtor is free(NULL) (payload 0): native + no-op, so
-   *  mlx may run it from the Metal completion thread without touching the
-   *  JS runtime (the old JSCallback dtor deadlocked serving / SIGTRAPed
-   *  under GC when that thread held the last reference — 2026-07-06). */
-  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, freeFnAddr,
-    );
-    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),
-    );
-  }
-
-  /** Graph-free LEAF with this array's exact bytes — ONE copy, straight from
-   *  the (evaluated) mlx buffer into a fresh mlx-owned array, no JS-heap
-   *  intermediate (`fromBytesCopy(rawBytes())` is TWO copies: device→JS heap
-   *  →new leaf). mlx_array_new_data copies synchronously, so the source only
-   *  has to outlive this call. The buffer is read LINEARLY — the caller must
-   *  pass a row-major contiguous array (ops.contiguous first for views), same
-   *  contract as rawBytes. Segmented-backward boundary detach hot path. */
-  detachCopy(): MlxArray {
-    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(`detachCopy: unsupported dtype ${this.dtypeName}`);
-    const sb = shapeBuf(this.shape);
-    return new MlxArray(C.mlx_array_new_data(p as never, ptr(sb), this.ndim, dt));
-  }
-
-  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 (lab/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 a supported floating-point or integer array (copy). */
-  rawBytes(): Uint8Array {
-    return this.rawBytesView().slice();
-  }
-
-  /** ZERO-COPY view of the evaluated array's bytes — aliases the mlx
-   *  buffer directly, no JS-heap copy. Valid only while THIS array is
-   *  alive and unmutated: do not retain past dispose(). Built for
-   *  bulk-write paths (kv-store's streamed writer) where per-tensor
-   *  rawBytes() copies pile up in the JS heap faster than GC collects
-   *  them — the residual A7 whole-entry RSS spike (2026-07-07). The
-   *  array must be contiguous (ops.contiguous first for views), same
-   *  contract as rawBytes. */
-  rawBytesView(): 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)
-      : dt === Dtype.int32 ? C.mlx_array_data_int32(this.handle)
-      : dt === Dtype.uint8 ? C.mlx_array_data_uint8(this.handle)
-      : dt === Dtype.int8 ? C.mlx_array_data_int8(this.handle)
-      : null;
-    if (p === null) throw new Error(`rawBytesView: unsupported dtype ${this.dtypeName}`);
-    return new Uint8Array(toArrayBuffer(p!, 0, this.nbytes));
-  }
-
-  /** Read back an INTEGER array without enqueueing a cast kernel: eval +
-   *  direct buffer read for uint32/int32. A cast (`astype`) is a NEW op that
-   *  queues BEHIND everything already dispatched on the GPU stream — on the
-   *  batched decode hot path, reading the pipeline register via toFloat32
-   *  stalled the "overlapped" token read for a FULL step of GPU work every
-   *  token (the core of the Phase-0 B=1 gap; unified-engine plan Phase 2).
-   *  Non-integer dtypes fall back to the rounded float readback. */
-  toIntTokens(): number[] {
-    if (this.dtype === Dtype.uint32) {
-      this.eval();
-      const p = C.mlx_array_data_uint32(this.handle);
-      return [...new Uint32Array(toArrayBuffer(p!, 0, this.size * 4))];
-    }
-    if (this.dtype === Dtype.int32) {
-      this.eval();
-      const p = C.mlx_array_data_int32(this.handle);
-      return [...new Int32Array(toArrayBuffer(p!, 0, this.size * 4))];
-    }
-    return [...this.toFloat32()].map((x) => Math.round(x));
-  }
-
-  /** 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;
-  }
-
-  /** Read an f32/f16/bf16 array numerically as float32 without enqueueing a
-   *  GPU cast. This evaluates the ORIGINAL array, then expands 16-bit values
-   *  on the host. Use this for small readbacks after later GPU work has
-   *  already been dispatched: toFloat32() would append astype(float32) behind
-   *  that work and turn the readback into an accidental full-step barrier. */
-  toFloat32Host(): Float32Array {
-    this.eval();
-    const dt = this.dtype;
-    if (dt === Dtype.float32) {
-      const p = C.mlx_array_data_float32(this.handle);
-      return new Float32Array(toArrayBuffer(p!, 0, this.size * 4)).slice();
-    }
-    if (dt === Dtype.bfloat16) {
-      const p = C.mlx_array_data_bfloat16(this.handle);
-      const src = new Uint16Array(toArrayBuffer(p!, 0, this.size * 2));
-      const words = new Uint32Array(src.length);
-      for (let i = 0; i < src.length; i++) words[i] = src[i]! << 16;
-      return new Float32Array(words.buffer);
-    }
-    if (dt === Dtype.float16) {
-      const p = C.mlx_array_data_float16(this.handle);
-      const src = new Uint16Array(toArrayBuffer(p!, 0, this.size * 2));
-      const out = new Float32Array(src.length);
-      for (let i = 0; i < src.length; i++) {
-        const bits = src[i]!;
-        const sign = bits & 0x8000;
-        const exponent = (bits >>> 10) & 0x1f;
-        const fraction = bits & 0x03ff;
-        let value: number;
-        if (exponent === 0) {
-          value = fraction === 0 ? 0 : fraction * 2 ** -24;
-        } else if (exponent === 0x1f) {
-          value = fraction === 0 ? Infinity : NaN;
-        } else {
-          value = (1 + fraction / 1024) * 2 ** (exponent - 15);
-        }
-        out[i] = sign ? -value : value;
-      }
-      return out;
-    }
-    throw new Error(`toFloat32Host: unsupported dtype ${this.dtypeName}`);
-  }
-
-  /** 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

@@ -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();
-  }
-}