@jax-js/jax 0.1.6 → 0.1.7

package/README.md

@@ -43,6 +43,55 @@ way to get started on a blank HTML page.
 </script>
 ```
 
+## Feature comparison
+
+Here's a quick, high-level comparison with other popular web ML runtimes:
+
+| Feature                         | jax-js     | TensorFlow.js   | onnxruntime-web    |
+| ------------------------------- | ---------- | --------------- | ------------------ |
+| **Overview**                    |            |                 |                    |
+| API style                       | JAX/NumPy  | TensorFlow-like | Static ONNX graphs |
+| Latest release                  | 2026       | ⚠️ 2024         | 2026               |
+| Speed                           | Fastest    | Fast            | Fastest            |
+| Bundle size (gzip)              | 80 KB      | 269 KB          | 90 KB + 24 MB Wasm |
+| **Autodiff & JIT**              |            |                 |                    |
+| Gradients                       | ✅         | ✅              | ❌                 |
+| Jacobian and Hessian            | ✅         | ❌              | ❌                 |
+| `jvp()` forward differentiation | ✅         | ❌              | ❌                 |
+| `jit()` kernel fusion           | ✅         | ❌              | ❌                 |
+| `vmap()` auto-vectorization     | ✅         | ❌              | ❌                 |
+| Graph capture                   | ✅         | ❌              | ✅                 |
+| **Backends & Data**             |            |                 |                    |
+| WebGPU backend                  | ✅         | 🟡 Preview      | ✅                 |
+| WebGL backend                   | ✅         | ✅              | ✅                 |
+| Wasm (CPU) backend              | ✅         | ✅              | ✅                 |
+| Eager array API                 | ✅         | ✅              | ❌                 |
+| Run ONNX models                 | 🟡 Partial | ❌              | ✅                 |
+| Read safetensors                | ✅         | ❌              | ❌                 |
+| Float64                         | ✅         | ❌              | ❌                 |
+| Float32                         | ✅         | ✅              | ✅                 |
+| Float16                         | ✅         | ❌              | ✅                 |
+| BFloat16                        | ❌         | ❌              | ❌                 |
+| Packed Uint8                    | ❌         | ❌              | 🟡 Partial         |
+| Mixed precision                 | ✅         | ❌              | ✅                 |
+| Mixed devices                   | ✅         | ❌              | ❌                 |
+| **Ops & Numerics**              |            |                 |                    |
+| Arithmetic functions            | ✅         | ✅              | ✅                 |
+| Matrix multiplication           | ✅         | ✅              | ✅                 |
+| General einsum                  | ✅         | 🟡 Partial      | 🟡 Partial         |
+| Sorting                         | ✅         | ❌              | ❌                 |
+| Activation functions            | ✅         | ✅              | ✅                 |
+| NaN/Inf numerics                | ✅         | ✅              | ✅                 |
+| Basic convolutions              | ✅         | ✅              | ✅                 |
+| n-d convolutions                | ✅         | ❌              | ✅                 |
+| Strided/dilated convolution     | ✅         | ✅              | ✅                 |
+| Cholesky, Lstsq                 | ✅         | ❌              | ❌                 |
+| LU, Solve, Determinant          | ✅         | ❌              | ❌                 |
+| SVD                             | ❌         | ❌              | ❌                 |
+| FFT                             | ✅         | ✅              | ✅                 |
+| Basic RNG (Uniform, Normal)     | ✅         | ✅              | ✅                 |
+| Advanced RNG                    | ✅         | ❌              | ❌                 |
+
 ## Tutorial
 
 Programming in `jax-js` looks [very similar to JAX](https://docs.jax.dev/en/latest/jax-101.html),
@@ -271,15 +320,18 @@ self-contained way in other projects.
 
 ### Performance
 
-We haven't spent a ton of time optimizing yet, but performance is generally pretty good. `jit` is
-very helpful for fusing operations together, and it's a feature only available on the web in jax-js.
-The default kernel-tuning heuristics get about 3000 GFLOP/s for matrix multiplication on an M4 Pro
-chip ([try it](https://jax-js.com/bench/matmul)).
+The WebGPU runtime includes an ML compiler with tile-aware optimizations, tuned for indiidual
+browsers. Also, this library uniquely has the `jit()` feature that fuses operations together and
+records an execution graph. jax-js achieves **over 7000 GFLOP/s** for matrix multiplication on an
+Apple M4 Max chip ([try it](https://jax-js.com/bench/matmul)).
 
-For that example, it's around the same GFLOP/s as
+For that example, it's significantly faster than both
 [TensorFlow.js](https://github.com/tensorflow/tfjs) and
 [ONNX Runtime Web](https://www.npmjs.com/package/onnxruntime-web), which both use handwritten
-libraries of custom kernels (versus jax-js, which generates kernels with an ML compiler).
+libraries of custom kernels.
+
+It's still early though. There's a lot of low-hanging fruit to continue optimizing the library, as
+well as unique optimizations such as FlashAttention variants.
 
 ### API Reference
 
@@ -291,8 +343,9 @@ That's all for this short tutorial. Please see the generated
 If you make something cool with jax-js, don't be a stranger! We can feature it here.
 
 - [Training neural networks on MNIST](https://jax-js.com/mnist)
+- [Voice cloning: Kyutai Pocket TTS](https://jax-js.com/tts)
 - [CLIP embeddings for books in-browser](https://jax-js.com/mobileclip)
-- [Object detection with DETR ResNet-50 (ONNX)](https://jax-js.com/detr-resnet-50)
+- [Object detection: DETR ResNet-50 (ONNX)](https://jax-js.com/detr-resnet-50)
 - [In-browser REPL](https://jax-js.com/repl)
 - [Matmul benchmark](https://jax-js.com/bench/matmul)
 - [Conv2d benchmark](https://jax-js.com/bench/conv2d)

package/dist/{backend-D7s-Retx.cjs → backend-B3foXiV_.cjs}

@@ -910,7 +910,7 @@ var AluExp = class AluExp {
 					return ret.simplify(cache);
 				}
 			}
-			if (y.arg > 0) {
+			if (y.arg > 0 && x.min >= 0) {
 				let [xNoConst, constVal] = [x, 0];
 				if (x.op === AluOp.Add && x.src[1].op === AluOp.Const) [xNoConst, constVal] = [x.src[0], x.src[1].arg];
 				const terms = [];
@@ -932,7 +932,7 @@ var AluExp = class AluExp {
 						rem = AluExp.add(rem, AluExp.mul(AluExp.const(x.dtype, Math.floor(remainders[i] / gcdVal)), terms[i]));
 						quo = AluExp.add(quo, AluExp.mul(AluExp.const(x.dtype, quotients[i]), terms[i]));
 					}
-					if (!((x.min < 0 || rem.min < 0) && remainders.some((r) => r !== 0))) if (op === AluOp.Mod) return AluExp.add(AluExp.mul(AluExp.const(x.dtype, gcdVal), AluExp.mod(rem, AluExp.const(x.dtype, Math.floor(y.arg / gcdVal)))), AluExp.const(x.dtype, constVal % gcdVal)).simplify(cache);
+					if (rem.min >= 0) if (op === AluOp.Mod) return AluExp.add(AluExp.mul(AluExp.const(x.dtype, gcdVal), AluExp.mod(rem, AluExp.const(x.dtype, Math.floor(y.arg / gcdVal)))), AluExp.const(x.dtype, constVal % gcdVal)).simplify(cache);
 					else return AluExp.add(AluExp.idiv(rem, AluExp.const(x.dtype, Math.floor(y.arg / gcdVal))), quo).simplify(cache);
 				}
 			}
@@ -4253,7 +4253,7 @@ async function createBackend(device) {
 		if (!navigator.gpu) return null;
 		const adapter = await navigator.gpu.requestAdapter({ powerPreference: "high-performance" });
 		if (!adapter) return null;
-		const { WebGPUBackend } = await Promise.resolve().then(() => require("./webgpu-rraa6dfz.cjs"));
+		const { WebGPUBackend } = await Promise.resolve().then(() => require("./webgpu-BykvF26B.cjs"));
 		const importantLimits = [
 			"maxBufferSize",
 			"maxComputeInvocationsPerWorkgroup",
@@ -4291,7 +4291,7 @@ async function createBackend(device) {
 		});
 		if (!gl) return null;
 		if (!gl.getExtension("EXT_color_buffer_float")) return null;
-		const { WebGLBackend } = await Promise.resolve().then(() => require("./webgl-CyfzNW8T.cjs"));
+		const { WebGLBackend } = await Promise.resolve().then(() => require("./webgl-DIIbKJ0G.cjs"));
 		return new WebGLBackend(gl);
 	} else throw new Error(`Backend not found: ${device}`);
 }

package/dist/{backend-Dx6Ob2D1.js → backend-nEolvdLv.js}

@@ -909,7 +909,7 @@ var AluExp = class AluExp {
 					return ret.simplify(cache);
 				}
 			}
-			if (y.arg > 0) {
+			if (y.arg > 0 && x.min >= 0) {
 				let [xNoConst, constVal] = [x, 0];
 				if (x.op === AluOp.Add && x.src[1].op === AluOp.Const) [xNoConst, constVal] = [x.src[0], x.src[1].arg];
 				const terms = [];
@@ -931,7 +931,7 @@ var AluExp = class AluExp {
 						rem = AluExp.add(rem, AluExp.mul(AluExp.const(x.dtype, Math.floor(remainders[i] / gcdVal)), terms[i]));
 						quo = AluExp.add(quo, AluExp.mul(AluExp.const(x.dtype, quotients[i]), terms[i]));
 					}
-					if (!((x.min < 0 || rem.min < 0) && remainders.some((r) => r !== 0))) if (op === AluOp.Mod) return AluExp.add(AluExp.mul(AluExp.const(x.dtype, gcdVal), AluExp.mod(rem, AluExp.const(x.dtype, Math.floor(y.arg / gcdVal)))), AluExp.const(x.dtype, constVal % gcdVal)).simplify(cache);
+					if (rem.min >= 0) if (op === AluOp.Mod) return AluExp.add(AluExp.mul(AluExp.const(x.dtype, gcdVal), AluExp.mod(rem, AluExp.const(x.dtype, Math.floor(y.arg / gcdVal)))), AluExp.const(x.dtype, constVal % gcdVal)).simplify(cache);
 					else return AluExp.add(AluExp.idiv(rem, AluExp.const(x.dtype, Math.floor(y.arg / gcdVal))), quo).simplify(cache);
 				}
 			}
@@ -4252,7 +4252,7 @@ async function createBackend(device) {
 		if (!navigator.gpu) return null;
 		const adapter = await navigator.gpu.requestAdapter({ powerPreference: "high-performance" });
 		if (!adapter) return null;
-		const { WebGPUBackend } = await import("./webgpu-C-VfevQW.js");
+		const { WebGPUBackend } = await import("./webgpu-B96vzWGE.js");
 		const importantLimits = [
 			"maxBufferSize",
 			"maxComputeInvocationsPerWorkgroup",
@@ -4290,7 +4290,7 @@ async function createBackend(device) {
 		});
 		if (!gl) return null;
 		if (!gl.getExtension("EXT_color_buffer_float")) return null;
-		const { WebGLBackend } = await import("./webgl-CLLvzJlO.js");
+		const { WebGLBackend } = await import("./webgl-DweKSWEm.js");
 		return new WebGLBackend(gl);
 	} else throw new Error(`Backend not found: ${device}`);
 }

package/dist/index.cjs

@@ -30,7 +30,7 @@ var __toESM = (mod$1, isNodeMode, target) => (target = mod$1 != null ? __create(
 }) : target, mod$1));
 
 //#endregion
-const require_backend = require('./backend-D7s-Retx.cjs');
+const require_backend = require('./backend-B3foXiV_.cjs');
 
 //#region src/frontend/convolution.ts
 /**
@@ -337,11 +337,11 @@ function map(fn, tree, ...rest) {
 }
 /** Take a reference of every array in a tree. */
 function ref(tree) {
-	return map((x) => x.ref, tree);
+	return map((x) => x instanceof Tracer ? x.ref : x, tree);
 }
 /** Dispose every array in a tree. */
 function dispose(tree) {
-	if (tree) map((x) => x.dispose(), tree);
+	if (tree) map((x) => x instanceof Tracer ? x.dispose() : void 0, tree);
 }
 
 //#endregion
@@ -615,14 +615,20 @@ function shrink(x, slice) {
 }
 function pad$1(x, width) {
 	const nd = ndim$1(x);
-	if (typeof width === "number") width = [[width, width]];
-	else if (require_backend.isNumberPair(width)) width = [width];
-	else if (!Array.isArray(width) || !width.every(require_backend.isNumberPair)) throw new TypeError(`Invalid pad() type: ${JSON.stringify(width)}`);
-	if (width.length === 1) {
-		const [w0, w1] = width[0];
-		width = require_backend.rep(nd, () => [w0, w1]);
-	} else if (width.length !== nd) throw new Error(`Invalid pad(): expected ${nd} axes, got ${width.length}`);
-	return bind1(Primitive.Pad, [x], { width });
+	let w;
+	if (typeof width === "number") w = [[width, width]];
+	else if (require_backend.isNumberPair(width)) w = [width];
+	else if (!Array.isArray(width)) {
+		const indicesAndPairs = Object.entries(width);
+		w = require_backend.rep(nd, [0, 0]);
+		for (const [k, v] of indicesAndPairs) w[require_backend.checkAxis(parseInt(k), nd)] = v;
+	} else if (!width.every(require_backend.isNumberPair)) throw new TypeError(`Invalid pad() type: ${JSON.stringify(width)}`);
+	else w = width;
+	if (w.length === 1) {
+		const [w0, w1] = w[0];
+		w = require_backend.rep(nd, () => [w0, w1]);
+	} else if (w.length !== nd) throw new Error(`Invalid pad(): expected ${nd} axes, got ${w.length}`);
+	return bind1(Primitive.Pad, [x], { width: w });
 }
 function triangularSolve$1(a, b, { lower = false, unitDiagonal = false } = {}) {
 	const as = getShape(a);
@@ -798,6 +804,22 @@ var Tracer = class Tracer {
 		const result = reduce(this.astype(castDtype), require_backend.AluOp.Add, axis, opts);
 		return result.mul(1 / n).astype(originalDtype);
 	}
+	/** Minimum of the elements of the array along a given axis. */
+	min(axis = null, opts) {
+		return reduce(this, require_backend.AluOp.Min, axis, opts);
+	}
+	/** Maximum of the elements of the array along a given axis. */
+	max(axis = null, opts) {
+		return reduce(this, require_backend.AluOp.Max, axis, opts);
+	}
+	/** Test whether all array elements along a given axis evaluate to true. */
+	all(axis = null, opts) {
+		return this.astype(require_backend.DType.Bool).min(axis, opts);
+	}
+	/** Test whether any array element along a given axis evaluates to true. */
+	any(axis = null, opts) {
+		return this.astype(require_backend.DType.Bool).max(axis, opts);
+	}
 	/** Permute the dimensions of an array. Defaults to reversing the axis order. */
 	transpose(perm) {
 		return transpose$1(this, perm);
@@ -5607,6 +5629,7 @@ __export(numpy_exports, {
 	moveaxis: () => moveaxis$1,
 	multiply: () => multiply,
 	nan: () => nan,
+	nanToNum: () => nanToNum,
 	ndim: () => ndim,
 	negative: () => negative,
 	notEqual: () => notEqual,
@@ -5804,24 +5827,22 @@ function max(a, axis = null, opts) {
 	return reduce(a, require_backend.AluOp.Max, axis, opts);
 }
 /**
-* Test whether all array elements along a given axis evaluate to True.
+* Test whether any array element along a given axis evaluates to True.
 *
 * Returns a boolean array with the same shape as `a` with the specified axis
 * removed. If axis is None, returns a scalar.
 */
-function all(a, axis = null, opts) {
-	a = fudgeArray(a).astype(require_backend.DType.Bool);
-	return min(a, axis, opts);
+function any(a, axis = null, opts) {
+	return fudgeArray(a).any(axis, opts);
 }
 /**
-* Test whether any array element along a given axis evaluates to True.
+* Test whether all array elements along a given axis evaluate to True.
 *
 * Returns a boolean array with the same shape as `a` with the specified axis
 * removed. If axis is None, returns a scalar.
 */
-function any(a, axis = null, opts) {
-	a = fudgeArray(a).astype(require_backend.DType.Bool);
-	return max(a, axis, opts);
+function all(a, axis = null, opts) {
+	return fudgeArray(a).all(axis, opts);
 }
 /** Return the peak-to-peak range along a given axis (`max - min`). */
 function ptp(a, axis = null, opts) {
@@ -5922,7 +5943,7 @@ function split$1(a, indicesOrSections, axis = 0) {
 		const partSize = size$1 / indicesOrSections;
 		sizes = require_backend.rep(indicesOrSections, partSize);
 	} else {
-		const indices = indicesOrSections;
+		const indices = indicesOrSections.map((i) => i < 0 ? i + size$1 : i);
 		sizes = [indices[0]];
 		for (let i = 1; i < indices.length; i++) sizes.push(indices[i] - indices[i - 1]);
 		sizes.push(size$1 - indices[indices.length - 1]);
@@ -6870,6 +6891,21 @@ function isposinf(x) {
 	return require_backend.isFloatDtype(x.dtype) ? x.equal(Infinity) : fullLike$1(x, false);
 }
 /**
+* Replace NaN and infinite entries in an array.
+*
+* By default, NaNs are replaced with `0.0`, and infinities are are substituted
+* with the corresponding maximum or minimum finite values.
+*/
+function nanToNum(x, { nan: nan$1 = 0, posinf = null, neginf = null } = {}) {
+	x = fudgeArray(x);
+	x = where(isnan(x.ref), nan$1, x);
+	posinf ??= require_backend.isFloatDtype(x.dtype) ? finfo(x.dtype).max : iinfo(x.dtype).max;
+	neginf ??= require_backend.isFloatDtype(x.dtype) ? finfo(x.dtype).min : iinfo(x.dtype).min;
+	x = where(isposinf(x.ref), posinf, x);
+	x = where(isneginf(x.ref), neginf, x);
+	return x;
+}
+/**
 * @function
 * Test element-wise for finite values (not infinity or NaN).
 */
@@ -7578,8 +7614,6 @@ function oneHot(x, numClasses) {
 *   `[B, L, N, H]`, or `[L, N, H]` if `B` is omitted.
 */
 function dotProductAttention(query, key$1, value, opts = {}) {
-	if (opts.querySeqLengths !== void 0 || opts.keyValueSeqLengths !== void 0) throw new Error("Sequence length masking is not yet implemented");
-	if (opts.localWindowSize !== void 0) throw new Error("Local attention is not yet implemented");
 	query = fudgeArray(query);
 	key$1 = fudgeArray(key$1);
 	value = fudgeArray(value);
@@ -7617,6 +7651,38 @@ function dotProductAttention(query, key$1, value, opts = {}) {
 		const causalMask = tri(L, S, 0, { dtype: require_backend.DType.Bool });
 		scores = where(causalMask, scores, -Infinity);
 	}
+	if (opts.localWindowSize !== void 0) {
+		const [before, after] = typeof opts.localWindowSize === "number" ? [opts.localWindowSize, opts.localWindowSize] : opts.localWindowSize;
+		if (before < 0 || after < 0 || !Number.isInteger(before) || !Number.isInteger(after)) throw new Error(`dotProductAttention: localWindowSize values must be non-negative, got ${opts.localWindowSize}`);
+		const localMask = tri(L, S, after, { dtype: require_backend.DType.Bool }).mul(tri(L, S, -before - 1, { dtype: require_backend.DType.Bool }).notEqual(true));
+		scores = where(localMask, scores, -Infinity);
+	}
+	if (opts.querySeqLengths !== void 0) {
+		const sl = expandDims(opts.querySeqLengths, [
+			-1,
+			-2,
+			-3
+		]);
+		scores = where(arange(L).reshape([
+			1,
+			1,
+			L,
+			1
+		]).less(sl), scores, -Infinity);
+	}
+	if (opts.keyValueSeqLengths !== void 0) {
+		const sl = expandDims(opts.keyValueSeqLengths, [
+			-1,
+			-2,
+			-3
+		]);
+		scores = where(arange(S).reshape([
+			1,
+			1,
+			1,
+			S
+		]).less(sl), scores, -Infinity);
+	}
 	const attn = softmax(scores, -1);
 	const out = einsum("BNLS,BSNH->BLNH", attn, value);
 	return isRank3 ? out.reshape([