mlx-bun 0.4.0 → 1.0.0

package/src/mlx/ops.ts

@@ -1,1021 +0,0 @@
-// Functional op layer over MlxArray. Parity-critical conventions:
-// - Python-float scalars in mlx promote weakly (bf16 array × float →
-//   bf16, scalar cast to bf16 first). `scalarLike` replicates that.
-// - Op composition order mirrors mlx-lm exactly; see model/gemma4.ts.
-
-import { ptr, read } from "bun:ffi";
-import { C, Dtype, type MlxHandle, optFloat, optInt, outArray, takeMlxError } from "./ffi";
-import { MlxArray, gpuStream } from "./array";
-
-const cstrCache = new Map<string, Buffer>();
-function cstr(s: string): Buffer {
-  let b = cstrCache.get(s);
-  if (!b) {
-    b = Buffer.from(s + "\0", "utf8");
-    cstrCache.set(s, b);
-  }
-  return b;
-}
-
-type S = MlxHandle;
-
-export function add(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("add", (o) => C.mlx_add(o, a.handle, b.handle, s)));
-}
-
-export function sub(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("subtract", (o) => C.mlx_subtract(o, a.handle, b.handle, s)));
-}
-
-export function mul(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("multiply", (o) => C.mlx_multiply(o, a.handle, b.handle, s)));
-}
-
-export function div(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("divide", (o) => C.mlx_divide(o, a.handle, b.handle, s)));
-}
-
-export function tanh(a: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("tanh", (o) => C.mlx_tanh(o, a.handle, s)));
-}
-
-export function sigmoid(a: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("sigmoid", (o) => C.mlx_sigmoid(o, a.handle, s)));
-}
-
-/** f32 → bf16 truncation with round-to-nearest-even, returning the 2
- *  little-endian bytes of the bf16 value. Matches mlx's
- *  `static_cast<bfloat16_t>(float)` bit-for-bit so a host-built bf16 scalar is
- *  numerically identical to `f32_scalar.astype(bf16)` — just without the GPU
- *  Copy (astype) kernel. */
-function bf16Bytes(value: number): Uint8Array {
-  const fb = new Float32Array([value]);
-  const bits = new Uint32Array(fb.buffer)[0]!;
-  let hi16: number;
-  if (((bits >>> 23) & 0xff) === 0xff) {
-    // inf / nan: keep the high 16 bits, and keep a NaN a (quiet) NaN.
-    hi16 = (bits >>> 16) & 0xffff;
-    if ((bits & 0x7fffff) !== 0) hi16 |= 0x40;
-  } else {
-    const roundingBias = 0x7fff + ((bits >>> 16) & 1); // round to nearest even
-    hi16 = ((bits + roundingBias) >>> 16) & 0xffff;
-  }
-  return new Uint8Array([hi16 & 0xff, (hi16 >>> 8) & 0xff]);
-}
-
-/** Scalar constant with the same dtype as `like` (python weak-scalar semantics).
- *  Built HOST-SIDE at the target dtype — mlx's own scalar→array path constructs
- *  a bf16/f32 constant directly, so `scalar * bf16_array` is a single multiply.
- *  Going through an f32 leaf + GPU `.astype(bf16)` (the old path) emits a
- *  separate Copy kernel (`v_copy…`) that the reference never dispatches. */
-export function scalarLike(value: number, like: MlxArray): MlxArray {
-  if (like.dtype === Dtype.float32)
-    return MlxArray.fromFloat32(new Float32Array([value]), []);
-  if (like.dtype === Dtype.bfloat16)
-    return MlxArray.fromBytesCopy(bf16Bytes(value), [], Dtype.bfloat16);
-  // f16 / other float dtypes: keep the astype path (rare here; no bit-exact JS
-  // converter written for them yet).
-  const f = MlxArray.fromFloat32(new Float32Array([value]), []);
-  const cast = f.astype(like.dtype);
-  f.dispose();
-  return cast;
-}
-
-export function mulScalar(a: MlxArray, value: number, s: S = gpuStream): MlxArray {
-  const k = scalarLike(value, a);
-  const r = mul(a, k, s);
-  k.dispose();
-  return r;
-}
-
-export function rmsNorm(
-  x: MlxArray, weight: MlxArray | null, eps: number, s: S = gpuStream,
-): MlxArray {
-  return new MlxArray(
-    outArray("fast_rms_norm", (o) =>
-      C.mlx_fast_rms_norm(o, x.handle, weight?.handle ?? 0n, eps, s),
-    ),
-  );
-}
-
-export function layerNorm(
-  x: MlxArray, weight: MlxArray | null, bias: MlxArray | null, eps: number,
-  s: S = gpuStream,
-): MlxArray {
-  return new MlxArray(
-    outArray("fast_layer_norm", (o) =>
-      C.mlx_fast_layer_norm(o, x.handle, weight?.handle ?? 0n, bias?.handle ?? 0n, eps, s),
-    ),
-  );
-}
-
-export function matmul(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("matmul", (o) => C.mlx_matmul(o, a.handle, b.handle, s)));
-}
-
-/** Bit-identical copy through a FRESH primitive node (mx.copy). Unlike
- *  `contiguous`, this always creates a new graph node even for an already-
- *  contiguous input — the tool for DETACHING an array from a multi-output
- *  primitive's sibling set (a retained sibling output pins every co-output's
- *  buffer natively; see gatedDeltaUpdate's state detach, 2026-08-20). */
-export function copyOf(a: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("copy", (o) => C.mlx_copy(o, a.handle, s)));
-}
-
-/** Materialize a row-major contiguous copy (needed before rawBytes on views). */
-export function contiguous(a: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(
-    outArray("contiguous", (o) => C.mlx_contiguous(o, a.handle, false, s)),
-  );
-}
-
-export function logicalOr(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(
-    outArray("logical_or", (o) => C.mlx_logical_or(o, a.handle, b.handle, s)),
-  );
-}
-
-export function rope(
-  x: MlxArray, dims: number, base: number | null, offset: number,
-  freqs: MlxArray | null, s: S = gpuStream,
-): MlxArray {
-  return new MlxArray(
-    outArray("fast_rope", (o) =>
-      C.mlx_fast_rope(
-        o, x.handle, dims, false, optFloat(base), 1.0, offset,
-        freqs?.handle ?? 0n, s,
-      ),
-    ),
-  );
-}
-
-/** fast::rope with every knob exposed (traditional + scale), for the
- *  universal rope factory (rope_utils.initialize_rope port). The plain
- *  `rope` above stays byte-identical for existing call sites. */
-export function ropeScaled(
-  x: MlxArray, dims: number, traditional: boolean, base: number | null,
-  scale: number, offset: number, freqs: MlxArray | null, s: S = gpuStream,
-): MlxArray {
-  return new MlxArray(
-    outArray("fast_rope", (o) =>
-      C.mlx_fast_rope(
-        o, x.handle, dims, traditional, optFloat(base), scale, offset,
-        freqs?.handle ?? 0n, s,
-      ),
-    ),
-  );
-}
-
-/** fast::rope with the position offset as an ARRAY (int32 scalar) instead
- *  of a baked int — required inside compiled decode graphs, where the
- *  offset changes every step but the graph must not. Same kernel as
- *  `rope`; bit-exactness vs the static form is asserted in
- *  tests/unit/compile.test.ts. */
-export function ropeDynamic(
-  x: MlxArray, dims: number, base: number | null, offset: MlxArray,
-  freqs: MlxArray | null, s: S = gpuStream,
-): MlxArray {
-  return new MlxArray(
-    outArray("fast_rope_dynamic", (o) =>
-      C.mlx_fast_rope_dynamic(
-        o, x.handle, dims, false, optFloat(base), 1.0, offset.handle,
-        freqs?.handle ?? 0n, s,
-      ),
-    ),
-  );
-}
-
-/** Array-offset fast::rope with every knob exposed (traditional + scale) —
- *  the dynamic twin of `ropeScaled`, for the universal rope factory's
- *  batched-decode path (per-row [B] offsets under left-padding). Same
- *  kernel as `ropeDynamic`; static/dynamic bit-exactness is asserted in
- *  tests/unit/compile.test.ts. */
-export function ropeScaledDynamic(
-  x: MlxArray, dims: number, traditional: boolean, base: number | null,
-  scale: number, offset: MlxArray, freqs: MlxArray | null, s: S = gpuStream,
-): MlxArray {
-  return new MlxArray(
-    outArray("fast_rope_dynamic", (o) =>
-      C.mlx_fast_rope_dynamic(
-        o, x.handle, dims, traditional, optFloat(base), scale, offset.handle,
-        freqs?.handle ?? 0n, s,
-      ),
-    ),
-  );
-}
-
-export function sdpa(
-  q: MlxArray, k: MlxArray, v: MlxArray, scale: number,
-  maskMode: "" | "causal" | "array", maskArr: MlxArray | null = null,
-  s: S = gpuStream,
-): MlxArray {
-  return new MlxArray(
-    outArray("fast_sdpa", (o) =>
-      C.mlx_fast_scaled_dot_product_attention(
-        o, q.handle, k.handle, v.handle, scale, ptr(cstr(maskMode)),
-        maskArr?.handle ?? 0n, 0n, false, s,
-      ),
-    ),
-  );
-}
-
-export interface QuantSpec {
-  bits: number;
-  groupSize: number;
-  mode: string;
-}
-
-export function quantizedMatmul(
-  x: MlxArray, w: MlxArray, scales: MlxArray, biases: MlxArray | null,
-  spec: QuantSpec, transpose = true, s: S = gpuStream,
-): MlxArray {
-  // UPSTREAM MLX BUG WORKAROUND (found 2026-07-02, reproduced on stock mlx
-  // 0.31.2): quantized_matmul with transpose=false returns WRONG VALUES when
-  // the row count is 2 or 3 (relnorm ~0.3-1.1 vs the dequantized matmul);
-  // M=1 and M>=4 are correct, transpose=true is correct at every M. Pad the
-  // row axis to 4 with zeros and slice the result back. NOTE this wrapper
-  // cannot protect mlx's INTERNAL vjp of a transpose=true qmm (C++ calls the
-  // primitive directly) — training heads pad the span through the head
-  // instead (loss.ts logitsFromHiddenPadM).
-  const rowAxis = x.shape.length - 2;
-  const rows = x.shape[rowAxis]!;
-  if (!transpose && (rows === 2 || rows === 3)) {
-    const padShape = [...x.shape];
-    padShape[rowAxis] = 4 - rows;
-    const pad = zeros(padShape, x.dtype, s);
-    const xPad = concatAxis([x, pad], rowAxis, s);
-    pad.dispose();
-    const yPad = quantizedMatmul(xPad, w, scales, biases, spec, transpose, s);
-    xPad.dispose();
-    const start = MlxArray.fromInt32(new Int32Array([0]), [1]);
-    const outShape = [...yPad.shape];
-    outShape[rowAxis] = rows;
-    const y = sliceDynamic(yPad, start, [rowAxis], outShape, s);
-    start.dispose();
-    yPad.dispose();
-    return y;
-  }
-  return new MlxArray(
-    outArray("quantized_matmul", (o) =>
-      C.mlx_quantized_matmul(
-        o, x.handle, w.handle, scales.handle, biases?.handle ?? 0n,
-        transpose, optInt(spec.groupSize), optInt(spec.bits),
-        ptr(cstr(spec.mode)), s,
-      ),
-    ),
-  );
-}
-
-/** Gathered quantized matmul over stacked expert weights (MoE).
- *  x [..., 1, K] with rhs_indices [...] selecting the expert per row;
- *  port of mx.gather_qmm as QuantizedSwitchLinear uses it. */
-export function gatherQmm(
-  x: MlxArray, w: MlxArray, scales: MlxArray, biases: MlxArray | null,
-  rhsIndices: MlxArray, spec: QuantSpec, sortedIndices: boolean,
-  transpose = true, s: S = gpuStream,
-): MlxArray {
-  return new MlxArray(
-    outArray("gather_qmm", (o) =>
-      C.mlx_gather_qmm(
-        o, x.handle, w.handle, scales.handle, biases?.handle ?? 0n,
-        0n /* lhs_indices */, rhsIndices.handle, transpose,
-        optInt(spec.groupSize), optInt(spec.bits), ptr(cstr(spec.mode)),
-        sortedIndices, s,
-      ),
-    ),
-  );
-}
-
-export function floorDivide(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(
-    outArray("floor_divide", (o) => C.mlx_floor_divide(o, a.handle, b.handle, s)),
-  );
-}
-
-export function dequantize(
-  w: MlxArray, scales: MlxArray, biases: MlxArray | null, spec: QuantSpec,
-  s: S = gpuStream,
-): MlxArray {
-  return new MlxArray(
-    outArray("dequantize", (o) =>
-      C.mlx_dequantize(
-        o, w.handle, scales.handle, biases?.handle ?? 0n,
-        optInt(spec.groupSize), optInt(spec.bits), ptr(cstr(spec.mode)),
-        0n, optInt(null), s,
-      ),
-    ),
-  );
-}
-
-export function takeAxis(a: MlxArray, indices: MlxArray, axis: number, s: S = gpuStream): MlxArray {
-  return new MlxArray(
-    outArray("take_axis", (o) => C.mlx_take_axis(o, a.handle, indices.handle, axis, s)),
-  );
-}
-
-export function reshape(a: MlxArray, shape: number[], s: S = gpuStream): MlxArray {
-  // bun:ffi cannot take a pointer to an empty view; ndim=0 ignores the dummy.
-  const buf = new Int32Array(shape.length ? shape : [0]);
-  return new MlxArray(
-    outArray("reshape", (o) => C.mlx_reshape(o, a.handle, ptr(buf), BigInt(shape.length), s)),
-  );
-}
-
-export function transposeAxes(a: MlxArray, axes: number[], s: S = gpuStream): MlxArray {
-  const buf = new Int32Array(axes);
-  return new MlxArray(
-    outArray("transpose", (o) => C.mlx_transpose_axes(o, a.handle, ptr(buf), BigInt(axes.length), s)),
-  );
-}
-
-export function argmaxAxis(a: MlxArray, axis: number, s: S = gpuStream): MlxArray {
-  return new MlxArray(
-    outArray("argmax_axis", (o) => C.mlx_argmax_axis(o, a.handle, axis, false, s)),
-  );
-}
-
-export function concatAxis(arrays: MlxArray[], axis: number, s: S = gpuStream): MlxArray {
-  const handles = new BigUint64Array(arrays.map((a) => a.handle));
-  const vec = C.mlx_vector_array_new_data(ptr(handles), BigInt(arrays.length));
-  try {
-    return new MlxArray(
-      outArray("concatenate", (o) => C.mlx_concatenate_axis(o, vec, axis, s)),
-    );
-  } finally {
-    C.mlx_vector_array_free(vec);
-  }
-}
-
-export function pow(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("power", (o) => C.mlx_power(o, a.handle, b.handle, s)));
-}
-
-/** nn.gelu_approx: 0.5 * x * (1 + tanh(√(2/π) * (x + 0.044715 * x³))) — composed
- *  exactly as mlx's python source, INCLUDING operand order: x**3 is mx.power
- *  (NOT x·x·x — they round differently in bf16), and every scalar·array product
- *  puts the SCALAR FIRST (`0.5 * x`, not `x * 0.5`). Multiply is commutative so
- *  this is bit-identical, but it makes the compiled geglu graph's operand string
- *  (`Multiply(C, V)`) match the oracle's fused kernel exactly. Scalars promote
- *  weakly to x's dtype. */
-export function geluApprox(x: MlxArray, s: S = gpuStream): MlxArray {
-  // EXACT copy of mlx's graph for
-  //   0.5 * x * (1 + tanh(√(2/π) * (x + 0.044715 * x**3)))
-  // as dumped via mx.export_to_dot. mlx's `__rmul__`/`__radd__` canonicalize
-  // every `scalar OP array` to `array OP scalar`, so every commutative product/
-  // sum below is ARRAY-FIRST, constant-second (NOT scalar-first). Build order
-  // follows Python left-to-right (`0.5 * x` first). Same libmlx compiles this to
-  // the identical fused kernel.
-  const cHalf = scalarLike(0.5, x);
-  const halfx = mul(x, cHalf, s);                           // Multiply(x, 0.5)
-  const cThree = scalarLike(3, x);
-  const x3 = pow(x, cThree, s);                             // Power(x, 3)
-  const cCoef = scalarLike(0.044715, x);
-  const cx3 = mul(x3, cCoef, s);                            // Multiply(x**3, 0.044715)
-  const inner = add(x, cx3, s);                             // Add(x, cx3)
-  const cSqrt = scalarLike(Math.sqrt(2 / Math.PI), x);
-  const scaled = mul(inner, cSqrt, s);                      // Multiply(inner, √(2/π))
-  const t = tanh(scaled, s);
-  const cOne = scalarLike(1, x);
-  const t1 = add(t, cOne, s);                               // Add(tanh, 1)
-  const out = mul(halfx, t1, s);                            // Multiply(halfx, 1+tanh)
-  for (const a of [cHalf, halfx, cThree, x3, cCoef, cx3, inner, cSqrt, scaled, t, cOne, t1]) a.dispose();
-  return out;
-}
-
-export function erf(a: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("erf", (o) => C.mlx_erf(o, a.handle, s)));
-}
-
-/** nn.gelu (PRECISE, erf-based): x · (1 + erf(x/√2)) / 2 — composed exactly
- *  as mlx's python source (weak scalars promote to x's dtype). Used by the
- *  gemma-1 gated MLP and starcoder2's plain MLP; distinct from geluApprox
- *  (gelu_pytorch_tanh). */
-export function geluPrecise(x: MlxArray, s: S = gpuStream): MlxArray {
-  const sqrt2 = scalarLike(Math.sqrt(2), x);
-  const xs = div(x, sqrt2, s); // x / math.sqrt(2) — division, not mul-by-reciprocal
-  const e = erf(xs, s);
-  const one = scalarLike(1, x);
-  const e1 = add(one, e, s);
-  const xe = mul(x, e1, s);
-  const two = scalarLike(2, x);
-  const out = div(xe, two, s); // (...) / 2
-  for (const a of [sqrt2, xs, e, one, e1, xe, two]) a.dispose();
-  return out;
-}
-
-/** mx.addmm(c, a, b): a@b + c fused (α=β=1) — nn.Linear's bias path
- *  (`mx.addmm(bias, x, weight.T)`). */
-export function addmm(c: MlxArray, a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(
-    outArray("addmm", (o) => C.mlx_addmm(o, c.handle, a.handle, b.handle, 1.0, 1.0, s)),
-  );
-}
-
-export function exp(a: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("exp", (o) => C.mlx_exp(o, a.handle, s)));
-}
-
-export function cos(a: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("cos", (o) => C.mlx_cos(o, a.handle, s)));
-}
-
-export function sin(a: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("sin", (o) => C.mlx_sin(o, a.handle, s)));
-}
-
-export function neg(a: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("negative", (o) => C.mlx_negative(o, a.handle, s)));
-}
-
-export function argsortAxis(a: MlxArray, axis: number, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("argsort", (o) => C.mlx_argsort_axis(o, a.handle, axis, s)));
-}
-
-export function argpartitionAxis(a: MlxArray, kth: number, axis: number, s: S = gpuStream): MlxArray {
-  return new MlxArray(
-    outArray("argpartition", (o) => C.mlx_argpartition_axis(o, a.handle, kth, axis, s)),
-  );
-}
-
-export function takeAlongAxis(a: MlxArray, idx: MlxArray, axis: number, s: S = gpuStream): MlxArray {
-  return new MlxArray(
-    outArray("take_along_axis", (o) => C.mlx_take_along_axis(o, a.handle, idx.handle, axis, s)),
-  );
-}
-
-export function putAlongAxis(
-  a: MlxArray, idx: MlxArray, values: MlxArray, axis: number, s: S = gpuStream,
-): MlxArray {
-  return new MlxArray(
-    outArray("put_along_axis", (o) =>
-      C.mlx_put_along_axis(o, a.handle, idx.handle, values.handle, axis, s),
-    ),
-  );
-}
-
-export function cumsum(a: MlxArray, axis: number, s: S = gpuStream): MlxArray {
-  return new MlxArray(
-    outArray("cumsum", (o) => C.mlx_cumsum_axis(o, a.handle, axis, false, true, 0n, s)),
-  );
-}
-
-/** Cumulative maximum along `axis` (inclusive, forward) — mlx.core.cummax. */
-export function cummax(a: MlxArray, axis: number, s: S = gpuStream): MlxArray {
-  return new MlxArray(
-    outArray("cummax", (o) => C.mlx_cummax_axis(o, a.handle, axis, false, true, s)),
-  );
-}
-
-export function where(cond: MlxArray, x: MlxArray, y: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(
-    outArray("where", (o) => C.mlx_where(o, cond.handle, x.handle, y.handle, s)),
-  );
-}
-
-// arange is built host-side and uploaded. The original mlx_arange binding
-// broke once the calling path got JIT-optimized; a standalone repro
-// (lab/repro/bun-ffi-f64/) traced the root cause to a Bun bug: after DFG
-// tier-up, typed-array reads following a bun:ffi call return stale values
-// (the JIT eliminates the load across the native call). Not f64 marshaling
-// — args reach C intact. See PLAN.md Phase 4 findings. Reuse large ranges
-// across attention layers and sampling, but bound their retained storage:
-// causal masks request a new range on every round once context exceeds 64K.
-const arangeCache = new Map<string, { array: MlxArray; bytes: number }>();
-const ARANGE_CACHE_MIN = 65536;
-const ARANGE_CACHE_BYTES = 8 * 1024 * 1024;
-let arangeCacheBytes = 0;
-
-export function arange(start: number, stop: number, step: number, dtype: Dtype, s: S = gpuStream): MlxArray {
-  if (!Number.isInteger(start) || !Number.isInteger(step) || step === 0)
-    throw new Error(`arange: integer start/step required (${start}, ${stop}, ${step})`);
-  const n = Math.max(0, Math.ceil((stop - start) / step));
-  const key = `${start}|${step}|${n}|${dtype}`;
-  const cached = arangeCache.get(key);
-  if (cached) {
-    arangeCache.delete(key);
-    arangeCache.set(key, cached);
-    return cached.array.slice([0], [n], s);
-  }
-
-  let arr: MlxArray;
-  if (dtype === Dtype.int32) {
-    // host-built int32 leaf (the Bun-f64-FFI-safe path — see note above).
-    const data = new Int32Array(n);
-    for (let i = 0; i < n; i++) data[i] = start + i * step;
-    arr = MlxArray.fromInt32(data, [n]);
-  } else {
-    // Native mlx Arange in the target dtype — matches mlx's `Arange` op exactly.
-    // The old `fromInt32 → .astype(dtype)` emitted a phantom int32→float Copy
-    // (an AsType the oracle never dispatches — e.g. the proportional-RoPE freq
-    // exponents). Integer start/step passed as doubles are exact.
-    arr = new MlxArray(outArray("arange", (o) => C.mlx_arange(o, start, stop, step, dtype, s)));
-  }
-  if (n >= ARANGE_CACHE_MIN) {
-    const bytes = arr.nbytes;
-    if (bytes <= ARANGE_CACHE_BYTES) {
-      while (arangeCacheBytes + bytes > ARANGE_CACHE_BYTES) {
-        const [oldKey, old] = arangeCache.entries().next().value!;
-        arangeCache.delete(oldKey);
-        arangeCacheBytes -= old.bytes;
-        // Returned slices retain their own native reference, including lazy
-        // views that have not been evaluated when this owner is evicted.
-        old.array.dispose();
-      }
-      arangeCache.set(key, { array: arr, bytes });
-      arangeCacheBytes += bytes;
-      return arr.slice([0], [n], s);
-    }
-  }
-  return arr;
-}
-
-export function logsumexpAxis(a: MlxArray, axis: number, keepdims: boolean, s: S = gpuStream): MlxArray {
-  return new MlxArray(
-    outArray("logsumexp", (o) => C.mlx_logsumexp_axis(o, a.handle, axis, keepdims, s)),
-  );
-}
-
-export function greaterEqual(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(
-    outArray("greater_equal", (o) => C.mlx_greater_equal(o, a.handle, b.handle, s)),
-  );
-}
-
-export function less(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("less", (o) => C.mlx_less(o, a.handle, b.handle, s)));
-}
-
-export function lessEqual(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("less_equal", (o) => C.mlx_less_equal(o, a.handle, b.handle, s)));
-}
-
-export function logicalNot(a: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("logical_not", (o) => C.mlx_logical_not(o, a.handle, s)));
-}
-
-export function equal(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("equal", (o) => C.mlx_equal(o, a.handle, b.handle, s)));
-}
-
-/** Reduce-all to a scalar (mlx.core.all). */
-export function allReduce(a: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("all", (o) => C.mlx_all(o, a.handle, false, s)));
-}
-
-/** any along one axis (mlx.core.any). */
-export function anyAxis(a: MlxArray, axis: number, keepdims: boolean, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("any_axis", (o) => C.mlx_any_axis(o, a.handle, axis, keepdims, s)));
-}
-
-/** Read a 0-d/1-elem boolean array to host (forces eval). */
-export function itemBool(a: MlxArray): boolean {
-  const u = a.astype(Dtype.uint32);
-  try {
-    return itemUint32(u) !== 0;
-  } finally {
-    u.dispose();
-  }
-}
-
-export function logicalAnd(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(
-    outArray("logical_and", (o) => C.mlx_logical_and(o, a.handle, b.handle, s)),
-  );
-}
-
-export function zeros(shape: number[], dtype: Dtype, s: S = gpuStream): MlxArray {
-  const buf = new Int32Array(shape);
-  return new MlxArray(
-    outArray("zeros", (o) => C.mlx_zeros(o, ptr(buf), BigInt(shape.length), dtype, s)),
-  );
-}
-
-export function sliceUpdate(
-  src: MlxArray, update: MlxArray, start: number[], stop: number[], s: S = gpuStream,
-): MlxArray {
-  const st = new Int32Array(start);
-  const sp = new Int32Array(stop);
-  const strides = new Int32Array(start.map(() => 1));
-  return new MlxArray(
-    outArray("slice_update", (o) =>
-      C.mlx_slice_update(
-        o, src.handle, update.handle,
-        ptr(st), BigInt(st.length), ptr(sp), BigInt(sp.length),
-        ptr(strides), BigInt(strides.length), s,
-      ),
-    ),
-  );
-}
-
-/** mx.slice with start as an ARRAY and the output shape as constants —
- *  the shapeless-compile-safe slice (mlx's Slice primitive cannot infer
- *  output shapes; DynamicSlice's output shape IS sliceSize). Used only
- *  inside compiled-decode traces for subrange slices whose bounds are
- *  per-layer constants (full-range slices get simplified away and need
- *  no substitute). Values identical to `slice`. */
-export function sliceDynamic(
-  a: MlxArray, start: MlxArray, axes: number[], sliceSize: number[], s: S = gpuStream,
-): MlxArray {
-  const ax = new Int32Array(axes);
-  const sz = new Int32Array(sliceSize);
-  return new MlxArray(
-    outArray("slice_dynamic", (o) =>
-      C.mlx_slice_dynamic(
-        o, a.handle, start.handle,
-        ptr(ax), BigInt(ax.length), ptr(sz), BigInt(sz.length), s,
-      ),
-    ),
-  );
-}
-
-/** slice_update with the start index as an ARRAY (one int32 per entry in
- *  `axes`) — the compiled-decode form of the per-step cache write. Same
- *  write as `sliceUpdate`; bit-exactness vs the static form is asserted
- *  in tests/unit/compile.test.ts. */
-export function sliceUpdateDynamic(
-  src: MlxArray, update: MlxArray, start: MlxArray, axes: number[], s: S = gpuStream,
-): MlxArray {
-  const ax = new Int32Array(axes);
-  return new MlxArray(
-    outArray("slice_update_dynamic", (o) =>
-      C.mlx_slice_update_dynamic(
-        o, src.handle, update.handle, start.handle,
-        ptr(ax), BigInt(ax.length), s,
-      ),
-    ),
-  );
-}
-
-export function randomKey(seed: bigint): MlxArray {
-  return new MlxArray(outArray("random_key", (o) => C.mlx_random_key(o, seed)));
-}
-
-/** Set the GLOBAL mlx PRNG seed (mlx.core.random.seed). Calls to randint/etc
- *  with a null key thread the global key, so seeding then calling in the same
- *  order as the reference reproduces its draws bit-for-bit. */
-export function randomSeed(seed: bigint): void {
-  if (C.mlx_random_seed(seed) !== 0) throw new Error(`mlx_random_seed failed: ${takeMlxError() ?? ""}`);
-}
-
-/** Uniform integers in [low, high) — mlx.core.random.randint. With key=null it
- *  uses (and advances) the global key, matching mx.random.randint exactly. */
-export function randint(
-  low: number, high: number, shape: number[], dtype: Dtype,
-  key: MlxArray | null = null, s: S = gpuStream,
-): MlxArray {
-  const buf = new Int32Array(shape);
-  const lo = MlxArray.fromInt32(new Int32Array([low]), []);
-  const hi = MlxArray.fromInt32(new Int32Array([high]), []);
-  try {
-    return new MlxArray(
-      outArray("random_randint", (o) =>
-        C.mlx_random_randint(o, lo.handle, hi.handle, ptr(buf), BigInt(shape.length), dtype, key?.handle ?? 0n, s),
-      ),
-    );
-  } finally {
-    lo.dispose();
-    hi.dispose();
-  }
-}
-
-export function randomCategorical(logits: MlxArray, key: MlxArray | null, s: S = gpuStream): MlxArray {
-  return new MlxArray(
-    outArray("categorical", (o) =>
-      C.mlx_random_categorical(o, logits.handle, -1, key?.handle ?? 0n, s),
-    ),
-  );
-}
-
-/** Eval several arrays in one call (e.g. cache state after a prefill chunk). */
-export function evalAll(arrays: MlxArray[]): void {
-  const handles = new BigUint64Array(arrays.map((a) => a.handle));
-  const vec = C.mlx_vector_array_new_data(ptr(handles), BigInt(arrays.length));
-  try {
-    if (C.mlx_eval(vec) !== 0) throw new Error(`mlx_eval failed: ${takeMlxError() ?? ""}`);
-  } finally {
-    C.mlx_vector_array_free(vec);
-  }
-}
-
-/** Kick off async evaluation (decode pipelining). */
-export function asyncEvalAll(arrays: MlxArray[]): void {
-  const handles = new BigUint64Array(arrays.map((a) => a.handle));
-  const vec = C.mlx_vector_array_new_data(ptr(handles), BigInt(arrays.length));
-  try {
-    if (C.mlx_async_eval(vec) !== 0) throw new Error(`mlx_async_eval failed: ${takeMlxError() ?? ""}`);
-  } finally {
-    C.mlx_vector_array_free(vec);
-  }
-}
-
-/** Read a scalar uint32 (forces eval). Hot per-token path: the out-param
- *  is read back via read.u32, never out[0] (DFG stale-read bug — see
- *  outArray in ffi.ts). */
-export function itemUint32(a: MlxArray): number {
-  a.eval();
-  const out = new Uint32Array(1);
-  const p = ptr(out);
-  if (C.mlx_array_item_uint32(p, a.handle) !== 0)
-    throw new Error("mlx_array_item_uint32 failed");
-  return read.u32(p, 0);
-}
-
-export function fromInt32(data: number[], shape: number[]): MlxArray {
-  return MlxArray.fromInt32(new Int32Array(data), shape);
-}
-
-// --- quantized KV support (Phase 6) ---------------------------------------
-
-export interface QuantizedTensor {
-  packed: MlxArray;
-  scales: MlxArray;
-  biases: MlxArray;
-}
-
-/** mx.quantize: w → (packed u32, scales, biases). `mode` defaults to the
- *  mlx-standard "affine" scheme; the quantizer seam may pass others. */
-export function quantize(
-  w: MlxArray, groupSize: number, bits: number, mode = "affine", s: S = gpuStream,
-): QuantizedTensor {
-  const slot = new BigUint64Array([C.mlx_vector_array_new()]);
-  const slotPtr = ptr(slot);
-  const status = C.mlx_quantize(
-    slotPtr, w.handle, optInt(groupSize), optInt(bits), ptr(cstr(mode)), 0n, s,
-  );
-  if (status !== 0) throw new Error(`quantize failed: ${takeMlxError() ?? ""}`);
-  const vec = read.u64(slotPtr, 0);
-  const get = (i: number): MlxArray => {
-    const aSlot = new BigUint64Array([C.mlx_array_new()]);
-    const aPtr = ptr(aSlot);
-    if (C.mlx_vector_array_get(aPtr, vec, BigInt(i)) !== 0)
-      throw new Error("vector_array_get failed");
-    return new MlxArray(read.u64(aPtr, 0));
-  };
-  try {
-    return { packed: get(0), scales: get(1), biases: get(2) };
-  } finally {
-    C.mlx_vector_array_free(vec);
-  }
-}
-
-// --- training: elementwise / reductions / random ------------------------
-
-/** Stops gradient flow through `a` (identity forward; ∂=0). DPO reference
- *  logprobs are stop_gradient'd before the policy forward. */
-export function stopGradient(a: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("stop_gradient", (o) => C.mlx_stop_gradient(o, a.handle, s)));
-}
-
-export function log(a: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("log", (o) => C.mlx_log(o, a.handle, s)));
-}
-
-export function logaddexp(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("logaddexp", (o) => C.mlx_logaddexp(o, a.handle, b.handle, s)));
-}
-
-export function sqrt(a: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("sqrt", (o) => C.mlx_sqrt(o, a.handle, s)));
-}
-
-export function rsqrt(a: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("rsqrt", (o) => C.mlx_rsqrt(o, a.handle, s)));
-}
-
-export function square(a: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("square", (o) => C.mlx_square(o, a.handle, s)));
-}
-
-export function abs(a: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("abs", (o) => C.mlx_abs(o, a.handle, s)));
-}
-
-/** Reduce-mean over all elements → scalar (keepdims=false) or [1,…]. */
-export function meanAll(a: MlxArray, keepdims = false, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("mean", (o) => C.mlx_mean(o, a.handle, keepdims, s)));
-}
-
-export function meanAxis(a: MlxArray, axis: number, keepdims: boolean, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("mean_axis", (o) => C.mlx_mean_axis(o, a.handle, axis, keepdims, s)));
-}
-
-export function clip(a: MlxArray, lo: MlxArray | null, hi: MlxArray | null, s: S = gpuStream): MlxArray {
-  return new MlxArray(
-    outArray("clip", (o) => C.mlx_clip(o, a.handle, lo?.handle ?? 0n, hi?.handle ?? 0n, s)),
-  );
-}
-
-/** mx.random.normal(shape, dtype, loc, scale, key). key may be null. */
-export function randomNormal(
-  shape: number[], dtype: Dtype, loc: number, scale: number,
-  key: MlxArray | null, s: S = gpuStream,
-): MlxArray {
-  const buf = new Int32Array(shape);
-  return new MlxArray(
-    outArray("random_normal", (o) =>
-      C.mlx_random_normal(o, ptr(buf), BigInt(shape.length), dtype, loc, scale, key?.handle ?? 0n, s),
-    ),
-  );
-}
-
-/** mx.random.uniform(low, high, shape, dtype, key). low/high are scalars.
- *  Matches mlx-lm LoRALinear A init: uniform(-1/√in, +1/√in). */
-export function randomUniform(
-  shape: number[], dtype: Dtype, low: number, high: number,
-  key: MlxArray | null, s: S = gpuStream,
-): MlxArray {
-  const buf = new Int32Array(shape);
-  const lo = MlxArray.fromFloat32(new Float32Array([low]), []);
-  const hi = MlxArray.fromFloat32(new Float32Array([high]), []);
-  try {
-    return new MlxArray(
-      outArray("random_uniform", (o) =>
-        C.mlx_random_uniform(o, lo.handle, hi.handle, ptr(buf), BigInt(shape.length), dtype, key?.handle ?? 0n, s),
-      ),
-    );
-  } finally {
-    lo.dispose();
-    hi.dispose();
-  }
-}
-
-/** mx.random.split(key, num) → [num, 2] uint32 subkeys. */
-export function randomSplitNum(key: MlxArray, num: number, s: S = gpuStream): MlxArray {
-  return new MlxArray(
-    outArray("random_split", (o) => C.mlx_random_split_num(o, key.handle, num, s)),
-  );
-}
-
-export function softmaxAxis(a: MlxArray, axis: number, precise: boolean, s: S = gpuStream): MlxArray {
-  return new MlxArray(
-    outArray("softmax", (o) => C.mlx_softmax_axis(o, a.handle, axis, precise, s)),
-  );
-}
-
-export function maxAxis(a: MlxArray, axis: number, keepdims: boolean, s: S = gpuStream): MlxArray {
-  return new MlxArray(
-    outArray("max_axis", (o) => C.mlx_max_axis(o, a.handle, axis, keepdims, s)),
-  );
-}
-
-export function sumAxis(a: MlxArray, axis: number, keepdims: boolean, s: S = gpuStream): MlxArray {
-  return new MlxArray(
-    outArray("sum_axis", (o) => C.mlx_sum_axis(o, a.handle, axis, keepdims, s)),
-  );
-}
-
-/** View-preserving expand_dims (reshape on a non-contiguous view copies —
- *  different kernel path, different reduction rounding; see Phase 6). */
-export function expandDims(a: MlxArray, axis: number, s: S = gpuStream): MlxArray {
-  return new MlxArray(
-    outArray("expand_dims", (o) => C.mlx_expand_dims(o, a.handle, axis, s)),
-  );
-}
-
-export function maximum(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(
-    outArray("maximum", (o) => C.mlx_maximum(o, a.handle, b.handle, s)),
-  );
-}
-
-/** quantized_matmul against a quantized tensor (x @ qt^T or x @ qt). */
-export function quantizedMatmulQT(
-  x: MlxArray, qt: QuantizedTensor, transpose: boolean,
-  groupSize: number, bits: number, s: S = gpuStream,
-): MlxArray {
-  return new MlxArray(
-    outArray("quantized_matmul", (o) =>
-      C.mlx_quantized_matmul(
-        o, x.handle, qt.packed.handle, qt.scales.handle, qt.biases.handle,
-        transpose, optInt(groupSize), optInt(bits), ptr(cstr("affine")), s,
-      ),
-    ),
-  );
-}
-
-/** nn.Conv1d: input [B, L, C_in], weight [C_out, K, C_in/groups] → [B, L', C_out].
- *  Depthwise (groups == channels) for the gated-DeltaNet causal conv; the model
- *  left-pads with the conv-state cache, so padding is 0 here. */
-export function conv1d(
-  input: MlxArray, weight: MlxArray, stride = 1, padding = 0, dilation = 1,
-  groups = 1, s: S = gpuStream,
-): MlxArray {
-  return new MlxArray(
-    outArray("conv1d", (o) =>
-      C.mlx_conv1d(o, input.handle, weight.handle, stride, padding, dilation, groups, s),
-    ),
-  );
-}
-
-/** nn.Conv2d: input [B, H, W, C_in] (NHWC), weight [C_out, kH, kW, C_in/groups]
- *  → [B, H', W', C_out]. mlx weight layout — the oracle transposes PyTorch's
- *  [C_out, C_in, kH, kW] at load. For the gemma-4 audio SSCP subsampler. */
-export function conv2d(
-  input: MlxArray, weight: MlxArray, stride: [number, number] = [1, 1],
-  padding: [number, number] = [0, 0], dilation: [number, number] = [1, 1],
-  groups = 1, s: S = gpuStream,
-): MlxArray {
-  return new MlxArray(
-    outArray("conv2d", (o) =>
-      C.mlx_conv2d(o, input.handle, weight.handle, stride[0], stride[1],
-        padding[0], padding[1], dilation[0], dilation[1], groups, s),
-    ),
-  );
-}
-
-/** nn.Conv3d: input [B, D, H, W, C_in] (NDHWC), weight [C_out, kD, kH, kW,
- *  C_in/groups] → [B, D', H', W', C_out]. The Qwen3-VL vision patch embed
- *  (kernel == stride) — the conv KERNEL's accumulation order is the parity-
- *  relevant arithmetic; a value-equivalent GEMM differs by bf16 ulp. */
-export function conv3d(
-  input: MlxArray, weight: MlxArray,
-  stride: [number, number, number] = [1, 1, 1],
-  padding: [number, number, number] = [0, 0, 0],
-  dilation: [number, number, number] = [1, 1, 1],
-  groups = 1, s: S = gpuStream,
-): MlxArray {
-  return new MlxArray(
-    outArray("conv3d", (o) =>
-      C.mlx_conv3d(o, input.handle, weight.handle, stride[0], stride[1],
-        stride[2], padding[0], padding[1], padding[2], dilation[0],
-        dilation[1], dilation[2], groups, s),
-    ),
-  );
-}
-
-/** mx.split(a, indices, axis): split at the given boundary offsets along `axis`
- *  into N+1 contiguous slices (mirrors numpy/mlx split-by-indices). */
-export function split(a: MlxArray, indices: number[], axis: number): MlxArray[] {
-  const shape = a.shape;
-  const ax = axis < 0 ? shape.length + axis : axis;
-  const dim = shape[ax]!;
-  const bounds = [0, ...indices, dim];
-  const out: MlxArray[] = [];
-  for (let i = 0; i < bounds.length - 1; i++) {
-    const start = shape.map(() => 0);
-    const stop = [...shape];
-    start[ax] = bounds[i]!;
-    stop[ax] = bounds[i + 1]!;
-    out.push(a.slice(start, stop));
-  }
-  return out;
-}
-
-/** nn.softplus: log(1 + exp(x)) == logaddexp(x, 0). */
-export function softplus(a: MlxArray, s: S = gpuStream): MlxArray {
-  const zero = scalarLike(0, a);
-  const r = logaddexp(a, zero, s);
-  zero.dispose();
-  return r;
-}
-
-/** nn.silu: x * sigmoid(x). */
-export function silu(a: MlxArray, s: S = gpuStream): MlxArray {
-  const sig = sigmoid(a, s);
-  const r = mul(a, sig, s);
-  sig.dispose();
-  return r;
-}
-
-/** fast::hadamard_transform over the last axis. `scale` defaults to
- *  1/sqrt(dim) (mlx's own default when unset) — pass it explicitly to
- *  match a reference that always states its normalization, per
- *  turboquant-ops.ts's fwht(). */
-export function hadamardTransform(x: MlxArray, scale: number | null = null, s: S = gpuStream): MlxArray {
-  return new MlxArray(
-    outArray("hadamard_transform", (o) =>
-      C.mlx_hadamard_transform(o, x.handle, optFloat(scale), s),
-    ),
-  );
-}
-
-export function minAxis(a: MlxArray, axis: number, keepdims: boolean, s: S = gpuStream): MlxArray {
-  return new MlxArray(
-    outArray("min_axis", (o) => C.mlx_min_axis(o, a.handle, axis, keepdims, s)),
-  );
-}
-
-/** mx.round(x, decimals=0) — round-half-to-even at the ULP grid mlx uses. */
-export function round(a: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("round", (o) => C.mlx_round(o, a.handle, 0, s)));
-}
-
-export function greater(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("greater", (o) => C.mlx_greater(o, a.handle, b.handle, s)));
-}
-
-export function bitwiseAnd(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("bitwise_and", (o) => C.mlx_bitwise_and(o, a.handle, b.handle, s)));
-}
-
-export function bitwiseOr(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("bitwise_or", (o) => C.mlx_bitwise_or(o, a.handle, b.handle, s)));
-}
-
-export function leftShift(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("left_shift", (o) => C.mlx_left_shift(o, a.handle, b.handle, s)));
-}
-
-export function rightShift(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
-  return new MlxArray(outArray("right_shift", (o) => C.mlx_right_shift(o, a.handle, b.handle, s)));
-}