mlx-bun 0.0.11 → 0.0.12

package/src/media-fetch.ts

@@ -13,10 +13,9 @@
 // apply.
 //
 // DNS names are resolved and every resolved address is checked before the
-// fetch. The fetch itself re-resolves, so a DNS-rebinding attacker who
-// controls TTLs can theoretically race the check — accepted for a local
-// single-user server (fixing it needs pinned-address dialing, which
-// bun's fetch doesn't expose).
+// fetch. The request then dials one of those checked IP literals while
+// preserving the logical Host header and HTTPS SNI, so the transport cannot
+// re-resolve to a private address between validation and connect.
 //
 // This module must stay free of mlx imports so tests/media-fetch.test.ts
 // runs model-free.
@@ -26,6 +25,8 @@ import { lookup } from "node:dns/promises";
 export interface MediaFetchPolicy {
   /** Permit private/loopback/link-local destinations (LAN hosts). */
   allowPrivate: boolean;
+  /** Optional caller-specific explanation when private destinations fail. */
+  privateDestinationHint?: string;
   /** Wall-clock budget for the whole fetch — every redirect hop plus the
    *  body read share one AbortSignal.timeout. */
   timeoutMs: number;
@@ -35,6 +36,39 @@ export interface MediaFetchPolicy {
   maxRedirects: number;
 }
 
+export type RestrictedFetch = (
+  input: string | URL,
+  init?: BunFetchRequestInit,
+) => Promise<Response>;
+
+export type ResolveHost = (
+  hostname: string,
+) => Promise<ReadonlyArray<{ address: string }>>;
+
+export interface RestrictedFetchDependencies {
+  fetch: RestrictedFetch;
+  resolve: ResolveHost;
+}
+
+export interface RestrictedFetchOptions {
+  /** Optional caller cancellation, combined with the policy timeout. */
+  signal?: AbortSignal;
+  /** Request headers applied on every validated redirect hop. */
+  headers?: Record<string, string>;
+  /** Test seam for deterministic DNS/redirect/body behavior. */
+  dependencies?: Partial<RestrictedFetchDependencies>;
+  /** Media callers reject non-2xx before reading; web_fetch returns the body. */
+  rejectHttpErrors?: boolean;
+}
+
+export interface RestrictedFetchResult {
+  ok: boolean;
+  status: number;
+  finalUrl: string;
+  contentType: string;
+  bytes: Uint8Array;
+}
+
 export function defaultMediaFetchPolicy(): MediaFetchPolicy {
   return {
     allowPrivate: process.env.MLX_BUN_ALLOW_PRIVATE_MEDIA === "1",
@@ -50,6 +84,10 @@ const PRIVATE_HINT =
   "private/loopback hosts are blocked by default; start the server with " +
   "--allow-private-media (or MLX_BUN_ALLOW_PRIVATE_MEDIA=1) to fetch from LAN hosts";
 
+function privateHint(policy: MediaFetchPolicy): string {
+  return policy.privateDestinationHint ?? PRIVATE_HINT;
+}
+
 function parseIpv4(s: string): number[] | null {
   const parts = s.split(".");
   if (parts.length !== 4) return null;
@@ -155,19 +193,22 @@ export function checkMediaUrl(u: URL, policy: MediaFetchPolicy): string | null {
   if (policy.allowPrivate) return null;
   const lower = host.toLowerCase();
   if (lower === "localhost" || lower.endsWith(".localhost"))
-    return `destination "${host}" — ${PRIVATE_HINT}`;
+    return `destination "${host}" — ${privateHint(policy)}`;
   // WHATWG URL parsing already canonicalized numeric IPv4 forms
   // (http://2130706433/, http://0x7f000001/ → 127.0.0.1).
   if (parseIpv4(lower) || parseIpv6(lower)) {
-    if (isBlockedAddress(lower)) return `destination "${host}" — ${PRIVATE_HINT}`;
+    if (isBlockedAddress(lower))
+      return `destination "${host}" — ${privateHint(policy)}`;
   }
   return null;
 }
 
-async function assertResolvesPublic(
+export async function assertResolvesPublic(
   hostname: string,
   kind: string,
   policy: MediaFetchPolicy,
+  resolve: ResolveHost = async (host) =>
+    await lookup(host, { all: true, verbatim: true }),
 ): Promise<void> {
   if (policy.allowPrivate) return;
   const host =
@@ -177,52 +218,105 @@ async function assertResolvesPublic(
   if (parseIpv4(host) || parseIpv6(host)) return; // literal — checked sync
   let addrs: { address: string }[];
   try {
-    addrs = await lookup(host, { all: true, verbatim: true });
+    addrs = [...await resolve(host)];
   } catch {
     throw new Error(`${kind} fetch failed: could not resolve host "${host}"`);
   }
   for (const a of addrs) {
     if (isBlockedAddress(a.address))
       throw new Error(
-        `${kind} url rejected: host "${host}" resolves to ${a.address} — ${PRIVATE_HINT}`,
+        `${kind} url rejected: host "${host}" resolves to ${a.address} — ${privateHint(policy)}`,
       );
   }
 }
 
-/** Fetch image/audio bytes for an OpenAI-style image_url/audio_url part.
- *  data: URLs decode locally (no policy — the bytes came in the request);
- *  http(s) URLs go through the destination policy above, with every
- *  redirect hop re-validated, one wall-clock timeout, and a streaming
- *  size cap. Throws Error with a client-presentable message — callers
- *  surface it as a 400 (the prompt-build catch in server.ts). */
-export async function fetchMediaBytes(
-  url: string,
-  kind: "image" | "audio",
-  policy: MediaFetchPolicy = defaultMediaFetchPolicy(),
-): Promise<Uint8Array> {
-  if (url.startsWith("data:")) {
-    const comma = url.indexOf(",");
-    if (comma === -1) throw new Error("malformed data: URL");
-    const meta = url.slice(0, comma);
-    const body = url.slice(comma + 1);
-    if (!meta.includes("base64")) throw new Error("data: URL must be base64");
-    return Uint8Array.from(Buffer.from(body, "base64"));
+async function resolveFetchAddresses(
+  hostname: string,
+  kind: string,
+  policy: MediaFetchPolicy,
+  resolve: ResolveHost,
+): Promise<string[]> {
+  const host =
+    hostname.startsWith("[") && hostname.endsWith("]")
+      ? hostname.slice(1, -1)
+      : hostname;
+  if (parseIpv4(host) || parseIpv6(host)) return [host];
+  if (policy.allowPrivate) return [];
+  let addrs: { address: string }[];
+  try {
+    addrs = [...await resolve(host)];
+  } catch {
+    throw new Error(`${kind} fetch failed: could not resolve host "${host}"`);
   }
+  if (addrs.length === 0)
+    throw new Error(`${kind} fetch failed: host "${host}" resolved to no addresses`);
+  for (const a of addrs) {
+    if (isBlockedAddress(a.address))
+      throw new Error(
+        `${kind} url rejected: host "${host}" resolves to ${a.address} — ${privateHint(policy)}`,
+      );
+  }
+  return addrs.map((a) => a.address);
+}
+
+/**
+ * Fetch an http(s) response through the shared public-destination policy.
+ * Every redirect is handled manually and revalidated, caller cancellation
+ * shares one signal with the wall-clock timeout, and the body is streamed
+ * into a byte-capped buffer.
+ */
+export async function fetchRestrictedHttpBytes(
+  url: string,
+  kind: string,
+  policy: MediaFetchPolicy,
+  options: RestrictedFetchOptions = {},
+): Promise<RestrictedFetchResult> {
   let current: URL;
   try {
     current = new URL(url);
   } catch {
     throw new Error(`unsupported ${kind} url scheme: ${url.slice(0, 16)}`);
   }
+
+  const fetchImpl = options.dependencies?.fetch ??
+    ((input: string | URL, init?: BunFetchRequestInit) => fetch(input, init));
+  const resolve = options.dependencies?.resolve ??
+    (async (host: string) => await lookup(host, { all: true, verbatim: true }));
+  const timeout = AbortSignal.timeout(policy.timeoutMs);
+  const signal = options.signal
+    ? AbortSignal.any([options.signal, timeout])
+    : timeout;
   const capMb = Math.round(policy.maxBytes / (1024 * 1024));
-  const signal = AbortSignal.timeout(policy.timeoutMs);
   let overCap = false;
+
   try {
     for (let hop = 0; ; hop++) {
       const reason = checkMediaUrl(current, policy);
       if (reason) throw new Error(`${kind} url rejected: ${reason}`);
-      await assertResolvesPublic(current.hostname, kind, policy);
-      const res = await fetch(current, { redirect: "manual", signal });
+      const addresses = await resolveFetchAddresses(
+        current.hostname,
+        kind,
+        policy,
+        resolve,
+      );
+      const logicalHostname = current.hostname.replace(/^\[|\]$/g, "");
+      const connectUrl = new URL(current);
+      if (!policy.allowPrivate && addresses.length > 0) {
+        const address = addresses[0]!;
+        connectUrl.hostname = address.includes(":") ? `[${address}]` : address;
+      }
+      const headers = new Headers(options.headers);
+      if (connectUrl.hostname !== current.hostname && !headers.has("host"))
+        headers.set("host", current.host);
+      const init: BunFetchRequestInit = {
+        redirect: "manual",
+        signal,
+        headers,
+      };
+      if (current.protocol === "https:" && connectUrl.hostname !== current.hostname) {
+        init.tls = { serverName: logicalHostname };
+      }
+      const res = await fetchImpl(connectUrl, init);
       if (res.status >= 300 && res.status < 400) {
         const loc = res.headers.get("location");
         if (!loc)
@@ -239,46 +333,94 @@ export async function fetchMediaBytes(
         }
         continue;
       }
-      if (!res.ok) throw new Error(`${kind} fetch failed: ${res.status} ${current}`);
+
+      if (options.rejectHttpErrors && !res.ok) {
+        try {
+          await res.body?.cancel();
+        } catch {}
+        throw new Error(`${kind} fetch failed: ${res.status} ${current}`);
+      }
+
       const declared = Number(res.headers.get("content-length") ?? "0");
-      if (declared > policy.maxBytes) {
+      if (Number.isFinite(declared) && declared > policy.maxBytes) {
         overCap = true;
         try {
           await res.body?.cancel();
         } catch {}
         throw new Error(`${kind} response exceeds the ${capMb} MB limit`);
       }
-      if (!res.body) return new Uint8Array(await res.arrayBuffer());
-      // stream + count: Content-Length is optional (chunked) and spoofable
-      const reader = res.body.getReader();
-      const chunks: Uint8Array[] = [];
-      let total = 0;
-      for (;;) {
-        const { done, value } = await reader.read();
-        if (done) break;
-        total += value.byteLength;
-        if (total > policy.maxBytes) {
+
+      let bytes: Uint8Array;
+      if (!res.body) {
+        bytes = new Uint8Array(await res.arrayBuffer());
+        if (bytes.byteLength > policy.maxBytes) {
           overCap = true;
-          try {
-            await reader.cancel();
-          } catch {}
           throw new Error(`${kind} response exceeds the ${capMb} MB limit`);
         }
-        chunks.push(value);
-      }
-      const out = new Uint8Array(total);
-      let off = 0;
-      for (const c of chunks) {
-        out.set(c, off);
-        off += c.byteLength;
+      } else {
+        const reader = res.body.getReader();
+        const chunks: Uint8Array[] = [];
+        let total = 0;
+        for (;;) {
+          const { done, value } = await reader.read();
+          if (done) break;
+          total += value.byteLength;
+          if (total > policy.maxBytes) {
+            overCap = true;
+            try {
+              await reader.cancel();
+            } catch {}
+            throw new Error(`${kind} response exceeds the ${capMb} MB limit`);
+          }
+          chunks.push(value);
+        }
+        bytes = new Uint8Array(total);
+        let offset = 0;
+        for (const chunk of chunks) {
+          bytes.set(chunk, offset);
+          offset += chunk.byteLength;
+        }
       }
-      return out;
+
+      return {
+        ok: res.ok,
+        status: res.status,
+        finalUrl: current.href,
+        contentType: res.headers.get("content-type") ?? "",
+        bytes,
+      };
     }
-  } catch (e) {
-    if (signal.aborted && !overCap)
+  } catch (error) {
+    if (timeout.aborted && !options.signal?.aborted && !overCap)
       throw new Error(
         `${kind} fetch timed out after ${policy.timeoutMs} ms: ${url.slice(0, 200)}`,
       );
-    throw e;
+    throw error;
+  }
+}
+
+/** Fetch image/audio bytes for an OpenAI-style image_url/audio_url part.
+ *  data: URLs decode locally (no policy — the bytes came in the request);
+ *  http(s) URLs go through the destination policy above, with every
+ *  redirect hop re-validated, one wall-clock timeout, and a streaming
+ *  size cap. Throws Error with a client-presentable message — callers
+ *  surface it as a 400 (the prompt-build catch in server.ts). */
+export async function fetchMediaBytes(
+  url: string,
+  kind: "image" | "audio",
+  policy: MediaFetchPolicy = defaultMediaFetchPolicy(),
+): Promise<Uint8Array> {
+  if (url.startsWith("data:")) {
+    const comma = url.indexOf(",");
+    if (comma === -1) throw new Error("malformed data: URL");
+    const meta = url.slice(0, comma);
+    const body = url.slice(comma + 1);
+    if (!meta.includes("base64")) throw new Error("data: URL must be base64");
+    return Uint8Array.from(Buffer.from(body, "base64"));
   }
+  return (
+    await fetchRestrictedHttpBytes(url, kind, policy, {
+      rejectHttpErrors: true,
+    })
+  ).bytes;
 }

package/src/mlx/array.ts

@@ -292,6 +292,49 @@ export class MlxArray {
     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);

package/src/mlx/autograd.ts

@@ -15,6 +15,23 @@ 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
@@ -28,18 +45,30 @@ import { MlxArray } from "./array";
  *  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 #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[]) {
+  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;
 
-    this.#cb = new JSCallback(
+    const cb = makeCallback(
       (outPtr: number, inVec: bigint, _payload: number): number => {
         try {
           // Read the primals out of the input vector_array. mlx_vector_array_get
@@ -73,16 +102,24 @@ export class ValueAndGrad {
       { 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");
+    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;
     }
-    this.#vag = read.u64(ptr(vagSlot), 0);
   }
 
   /** Run the loss + backward on `primals` (in argument order). Returns the
@@ -136,8 +173,8 @@ export class ValueAndGrad {
   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.#freeVag(this.#vag);
+    this.#freeClosure(this.#closure);
     this.#cb.close();
   }
 

package/src/mlx/checkpoint.ts

@@ -15,6 +15,22 @@ 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
@@ -24,14 +40,24 @@ import { MlxArray } from "./array";
  *  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 #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[]) {
-    this.#cb = new JSCallback(
+  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));
@@ -58,15 +84,22 @@ export class Checkpoint {
       { 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() ?? ""}`);
+    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;
     }
-    this.#ckpt = read.u64(ptr(slot), 0);
   }
 
   /** Apply the checkpointed fn. Returns its outputs as owned MlxArrays. */
@@ -105,8 +138,8 @@ export class Checkpoint {
   dispose(): void {
     if (this.#disposed) return;
     this.#disposed = true;
-    C.mlx_closure_free(this.#ckpt);
-    C.mlx_closure_free(this.#closure);
+    this.#freeClosure(this.#ckpt);
+    this.#freeClosure(this.#closure);
     this.#cb.close();
   }