mlx-bun 0.0.10 → 0.0.12

package/src/mlx/autograd.ts

@@ -1,6 +1,6 @@
 // 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:
+// 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.
 //
@@ -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,24 +45,36 @@ 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
           // 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).
+          // 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++) {
@@ -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();
   }
 

package/src/mlx/custom-vjp.ts

@@ -15,6 +15,24 @@ 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 CustomVjpConstructorHooks {
+  callback?: (
+    callback: ConstructorParameters<typeof JSCallback>[0],
+    definition: ConstructorParameters<typeof JSCallback>[1],
+  ) => CallbackHandle;
+  forwardClosureNew?: NativeCall<typeof C.mlx_closure_new_func_payload>;
+  vjpClosureNew?: NativeCall<typeof C.mlx_closure_custom_new_func_payload>;
+  transformedNew?: NativeCall<typeof C.mlx_closure_new>;
+  transform?: NativeCall<typeof C.mlx_custom_vjp>;
+  closureFree?: NativeCall<typeof C.mlx_closure_free>;
+  customClosureFree?: NativeCall<typeof C.mlx_closure_custom_free>;
+}
+
 function readVec(vec: bigint): MlxArray[] {
   const n = Number(C.mlx_vector_array_size(vec));
   const out: MlxArray[] = [];
@@ -41,19 +59,30 @@ function setVec(outPtr: number, arrays: MlxArray[]): void {
  *  pure functions of their arguments — the vjp re-reads everything from the
  *  graph during the enclosing backward. */
 export class CustomVjp {
-  readonly #fwdCb: JSCallback;
-  readonly #vjpCb: JSCallback;
+  readonly #fwdCb: CallbackHandle;
+  readonly #vjpCb: CallbackHandle;
   readonly #fwdClosure: bigint;
   readonly #vjpClosure: bigint;
   readonly #combined: bigint;
+  readonly #freeClosure: NativeCall<typeof C.mlx_closure_free>;
+  readonly #freeCustomClosure: NativeCall<typeof C.mlx_closure_custom_free>;
   #error: string | null = null;
   #disposed = false;
 
   constructor(
     fwd: (inputs: MlxArray[]) => MlxArray[],
     vjp: (primals: MlxArray[], cotangents: MlxArray[], outputs: MlxArray[]) => MlxArray[],
+    hooks: CustomVjpConstructorHooks = {},
   ) {
-    this.#fwdCb = new JSCallback(
+    const makeCallback = hooks.callback ?? ((callback, definition) => new JSCallback(callback, definition));
+    const forwardClosureNew = hooks.forwardClosureNew ?? C.mlx_closure_new_func_payload;
+    const vjpClosureNew = hooks.vjpClosureNew ?? C.mlx_closure_custom_new_func_payload;
+    const transformedNew = hooks.transformedNew ?? C.mlx_closure_new;
+    const transform = hooks.transform ?? C.mlx_custom_vjp;
+    this.#freeClosure = hooks.closureFree ?? C.mlx_closure_free;
+    this.#freeCustomClosure = hooks.customClosureFree ?? C.mlx_closure_custom_free;
+
+    const fwdCb = makeCallback(
       (outPtr: number, inVec: bigint, _payload: number): number => {
         try {
           const inputs = readVec(inVec);
@@ -70,37 +99,56 @@ export class CustomVjp {
       { args: ["ptr", "u64", "ptr"], returns: "i32" },
     );
 
-    this.#vjpCb = new JSCallback(
-      (resPtr: number, primalsVec: bigint, cotsVec: bigint, outsVec: bigint, _payload: number): number => {
-        try {
-          const primals = readVec(primalsVec);
-          const cots = readVec(cotsVec);
-          const outs = readVec(outsVec);
-          const grads = vjp(primals, cots, outs);
-          setVec(resPtr, grads);
-          for (const a of [...primals, ...cots, ...outs]) a.dispose();
-          for (const g of grads) g.dispose();
-          return 0;
-        } catch (e) {
-          this.#error = e instanceof Error ? e.message : String(e);
-          return 1;
-        }
-      },
-      { args: ["ptr", "u64", "u64", "u64", "ptr"], returns: "i32" },
-    );
-
-    this.#fwdClosure = C.mlx_closure_new_func_payload(this.#fwdCb.ptr as never, null, null);
-    this.#vjpClosure = C.mlx_closure_custom_new_func_payload(this.#vjpCb.ptr as never, null, null);
+    let vjpCb: CallbackHandle;
+    try {
+      vjpCb = makeCallback(
+        (resPtr: number, primalsVec: bigint, cotsVec: bigint, outsVec: bigint, _payload: number): number => {
+          try {
+            const primals = readVec(primalsVec);
+            const cots = readVec(cotsVec);
+            const outs = readVec(outsVec);
+            const grads = vjp(primals, cots, outs);
+            setVec(resPtr, grads);
+            for (const a of [...primals, ...cots, ...outs]) a.dispose();
+            for (const g of grads) g.dispose();
+            return 0;
+          } catch (e) {
+            this.#error = e instanceof Error ? e.message : String(e);
+            return 1;
+          }
+        },
+        { args: ["ptr", "u64", "u64", "u64", "ptr"], returns: "i32" },
+      );
+    } catch (e) {
+      // Callback construction is itself fallible. The native-closure guard
+      // below starts only after both callbacks exist, so this partial state
+      // must release the already-created forward callback here.
+      fwdCb.close();
+      throw e;
+    }
 
-    const slot = new BigUint64Array([C.mlx_closure_new()]);
-    if (C.mlx_custom_vjp(ptr(slot), this.#fwdClosure, this.#vjpClosure) !== 0) {
-      C.mlx_closure_free(this.#fwdClosure);
-      C.mlx_closure_custom_free(this.#vjpClosure);
-      this.#fwdCb.close();
-      this.#vjpCb.close();
-      throw new Error(`mlx_custom_vjp failed: ${takeMlxError() ?? ""}`);
+    let fwdClosure: bigint | null = null;
+    let vjpClosure: bigint | null = null;
+    let slot: BigUint64Array | null = null;
+    try {
+      fwdClosure = forwardClosureNew(fwdCb.ptr as never, null, null);
+      vjpClosure = vjpClosureNew(vjpCb.ptr as never, null, null);
+      slot = new BigUint64Array([transformedNew()]);
+      if (transform(ptr(slot), fwdClosure, vjpClosure) !== 0)
+        throw new Error(`mlx_custom_vjp failed: ${takeMlxError() ?? ""}`);
+      this.#fwdCb = fwdCb;
+      this.#vjpCb = vjpCb;
+      this.#fwdClosure = fwdClosure;
+      this.#vjpClosure = vjpClosure;
+      this.#combined = read.u64(ptr(slot), 0);
+    } catch (e) {
+      if (slot) this.#freeClosure(read.u64(ptr(slot), 0));
+      if (fwdClosure !== null) this.#freeClosure(fwdClosure);
+      if (vjpClosure !== null) this.#freeCustomClosure(vjpClosure);
+      fwdCb.close();
+      vjpCb.close();
+      throw e;
     }
-    this.#combined = read.u64(ptr(slot), 0);
   }
 
   apply(inputs: MlxArray[]): MlxArray[] {
@@ -127,9 +175,9 @@ export class CustomVjp {
   dispose(): void {
     if (this.#disposed) return;
     this.#disposed = true;
-    C.mlx_closure_free(this.#combined);
-    C.mlx_closure_free(this.#fwdClosure);
-    C.mlx_closure_custom_free(this.#vjpClosure);
+    this.#freeClosure(this.#combined);
+    this.#freeClosure(this.#fwdClosure);
+    this.#freeCustomClosure(this.#vjpClosure);
     this.#fwdCb.close();
     this.#vjpCb.close();
   }

package/src/mlx/ffi.ts

@@ -15,7 +15,7 @@ import { existsSync } from "node:fs";
 import { dirname, join } from "node:path";
 import { nativePackDir } from "../native-pack";
 
-const { ptr: P, i32, u64, f32, cstring } = FFIType;
+const { ptr: P, i32, u64, f32, f64, cstring } = FFIType;
 
 /** libmlxc resolution, in order: explicit env override → next to the
  *  executable (the `bun build --compile` sidecar layout, where
@@ -57,6 +57,9 @@ export const C = dlopen(LIBMLXC_PATH, {
   mlx_array_data_float16: { args: [u64], returns: P },
   mlx_array_data_bfloat16: { args: [u64], returns: P },
   mlx_array_data_uint32: { args: [u64], returns: P },
+  mlx_array_data_int32: { args: [u64], returns: P },
+  mlx_array_data_uint8: { args: [u64], returns: P },
+  mlx_array_data_int8: { args: [u64], returns: P },
   mlx_array_item_float32: { args: [P, u64], returns: i32 },
   // streams
   mlx_default_gpu_stream_new: { args: [], returns: u64 },
@@ -142,6 +145,7 @@ export const C = dlopen(LIBMLXC_PATH, {
   mlx_negative: { args: [P, u64, u64], returns: i32 },
   mlx_argsort_axis: { args: [P, u64, i32, u64], returns: i32 },
   mlx_argpartition_axis: { args: [P, u64, i32, i32, u64], returns: i32 },
+  mlx_arange: { args: [P, f64, f64, f64, i32, u64], returns: i32 },
   mlx_take_along_axis: { args: [P, u64, u64, i32, u64], returns: i32 },
   mlx_put_along_axis: { args: [P, u64, u64, u64, i32, u64], returns: i32 },
   mlx_cumsum: { args: [P, u64, i32, FFIType.bool, FFIType.bool, u64], returns: i32 },
@@ -169,7 +173,17 @@ export const C = dlopen(LIBMLXC_PATH, {
   // (res, input, weight, stride, padding, dilation, groups, stream) — depthwise
   // causal conv1d for Qwen3.5 gated-DeltaNet (groups == channels, padding 0).
   mlx_conv1d: { args: [P, u64, u64, i32, i32, i32, i32, u64], returns: i32 },
-  // --- training: autograd (value_and_grad) — proven in spikes/phase-train-vag.ts.
+  // C signature: (res, input, weight, stride_0, stride_1, padding_0,
+  // padding_1, dilation_0, dilation_1, groups, stream) — NHWC conv2d for the
+  // gemma-4 audio SSCP subsampler.
+  // HAZARD (lab/repro/bun-ffi-stack-args): Apple arm64 packs sub-8-byte STACK
+  // args at natural size (dilation_1@sp+0, groups@sp+4, stream@sp+8) but
+  // bun:ffi 1.3.14 writes one 8-byte slot per arg — the callee reads shifted
+  // garbage and mlx segfaults on a bogus stream. Workaround: declare the two
+  // stack ints as ONE u64 (dilation_1 | groups<<32) so Bun's slot layout
+  // byte-matches the Apple ABI. ops.conv2d does the packing.
+  mlx_conv2d: { args: [P, u64, u64, i32, i32, i32, i32, i32, u64, u64], returns: i32 },
+  // --- training: autograd (value_and_grad) — proven in lab/spikes/phase-train-vag.ts.
   // mlx_value_and_grad(res: mlx_closure_value_and_grad*, fun: mlx_closure,
   //   argnums: int*, n) → builds a value+grad closure differentiating the
   //   selected inputs of `fun` (a mlx_closure_new_func_payload over a JS loss).
@@ -260,6 +274,16 @@ export const C = dlopen(LIBMLXC_PATH, {
   // GPU trace capture (Phase E step 2: size the prize before the kernel)
   mlx_metal_start_capture: { args: [cstring], returns: i32 },
   mlx_metal_stop_capture: { args: [], returns: i32 },
+  // TurboQuant KV codec (Phase 13): FWHT rotation + per-group affine/Lloyd-Max
+  // (res, a, opt scale, stream) — 1/sqrt(dim) normalization when scale is null.
+  mlx_hadamard_transform: { args: [P, u64, u64, u64], returns: i32 },
+  mlx_min_axis: { args: [P, u64, i32, FFIType.bool, u64], returns: i32 },
+  mlx_round: { args: [P, u64, i32, u64], returns: i32 },
+  mlx_greater: { args: [P, u64, u64, u64], returns: i32 },
+  mlx_bitwise_and: { args: [P, u64, u64, u64], returns: i32 },
+  mlx_bitwise_or: { args: [P, u64, u64, u64], returns: i32 },
+  mlx_left_shift: { args: [P, u64, u64, u64], returns: i32 },
+  mlx_right_shift: { args: [P, u64, u64, u64], returns: i32 },
 }).symbols;
 
 export type MlxHandle = bigint;
@@ -329,7 +353,7 @@ export const NULL_HANDLE = 0n;
 // RULE: never read a typed array that a bun:ffi call wrote through a
 // pointer — once the calling function is DFG-compiled, the JIT eliminates
 // the load across the native call and `buf[i]` returns stale values
-// (repro/bun-ffi-f64/ISSUE.md; PLAN.md Phase 4 findings). Out-param slots
+// (lab/repro/bun-ffi-f64/ISSUE.md; PLAN.md Phase 4 findings). Out-param slots
 // must be read back with bun:ffi `read.*` (verified safe). Initializing the
 // slot via the typed-array *constructor* is fine: that store happens in
 // host code the JIT can't elide.

package/src/mlx/metal-kernel.ts

@@ -89,15 +89,18 @@ export class MetalKernel {
       C.mlx_fast_metal_kernel_config_set_thread_group(cfg, call.threadGroup[0], call.threadGroup[1], call.threadGroup[2]);
       if (call.initValue !== undefined)
         C.mlx_fast_metal_kernel_config_set_init_value(cfg, call.initValue);
-      for (const [k, v] of Object.entries(call.templateInts ?? {})) {
+      // Dtypes-first, then ints, to MATCH mlx's own specialization-name order
+      // (mlx-lm's gated_delta_step passes [("InT",…),("StT",…),("Dk",…),…]).
+      // Order affects only the generated name/cache-key, not binding (named args).
+      for (const [k, v] of Object.entries(call.templateDtypes ?? {})) {
         const kb = Buffer.from(k + "\0", "utf8");
         keepAlive.push(kb);
-        C.mlx_fast_metal_kernel_config_add_template_arg_int(cfg, ptr(kb), v);
+        C.mlx_fast_metal_kernel_config_add_template_arg_dtype(cfg, ptr(kb), v);
       }
-      for (const [k, v] of Object.entries(call.templateDtypes ?? {})) {
+      for (const [k, v] of Object.entries(call.templateInts ?? {})) {
         const kb = Buffer.from(k + "\0", "utf8");
         keepAlive.push(kb);
-        C.mlx_fast_metal_kernel_config_add_template_arg_dtype(cfg, ptr(kb), v);
+        C.mlx_fast_metal_kernel_config_add_template_arg_int(cfg, ptr(kb), v);
       }
 
       const handles = new BigUint64Array(inputs.map((a) => a.handle));