@jax-js/jax 0.1.5 → 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-DziQSaoQ.cjs → backend-B3foXiV_.cjs}

@@ -151,6 +151,19 @@ function normalizeAxis(axis, ndim) {
 		return sorted(seen);
 	}
 }
+/** Check for an array of integers with no duplicates. */
+function checkInts(indices) {
+	if (typeof indices === "number") {
+		if (!Number.isInteger(indices)) throw new TypeError(`Expected integer index, got ${indices}`);
+	} else {
+		const seen = /* @__PURE__ */ new Set();
+		for (const i of indices) {
+			if (!Number.isInteger(i)) throw new TypeError(`Expected integer indices, got ${i}`);
+			if (seen.has(i)) throw new Error(`Duplicate index ${i} passed to function`);
+			seen.add(i);
+		}
+	}
+}
 function range(start, stop, step = 1) {
 	if (stop === void 0) {
 		stop = start;
@@ -897,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 = [];
@@ -919,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);
 				}
 			}
@@ -2306,10 +2319,10 @@ function tuneWebgpu(kernel) {
 			upcastedAxis.add(choices[0][2]);
 		} else break;
 	}
-	if (/chrome/i.test(navigator.userAgent) && dim.reduce < dim.unroll && (prod(dim.st.shape.slice(dim.unroll)) <= 4 || dim.unroll === dim.upcast && prod(dim.st.shape.slice(dim.upcast)) < 64)) {
+	if (!/Mobi|Android/i.test(navigator.userAgent) && dim.reduce < dim.unroll && (prod(dim.st.shape.slice(dim.unroll)) <= 4 || dim.unroll === dim.upcast && prod(dim.st.shape.slice(dim.upcast)) < 64)) {
 		const s = dim.st.shape[dim.unroll - 1];
 		if (0 < s && s <= 32) dim.applyUnroll(dim.reduce, s);
-		else for (const splits of [4]) if (s % splits === 0) {
+		else for (const splits of [8, 4]) if (s % splits === 0) {
 			dim.applyUnroll(dim.unroll - 1, splits);
 			break;
 		}
@@ -4240,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-Db2JrNBr.cjs"));
+		const { WebGPUBackend } = await Promise.resolve().then(() => require("./webgpu-BykvF26B.cjs"));
 		const importantLimits = [
 			"maxBufferSize",
 			"maxComputeInvocationsPerWorkgroup",
@@ -4278,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-ClIYb8jP.cjs"));
+		const { WebGLBackend } = await Promise.resolve().then(() => require("./webgl-DIIbKJ0G.cjs"));
 		return new WebGLBackend(gl);
 	} else throw new Error(`Backend not found: ${device}`);
 }
@@ -4446,6 +4459,12 @@ Object.defineProperty(exports, 'checkAxis', {
     return checkAxis;
   }
 });
+Object.defineProperty(exports, 'checkInts', {
+  enumerable: true,
+  get: function () {
+    return checkInts;
+  }
+});
 Object.defineProperty(exports, 'deepEqual', {
   enumerable: true,
   get: function () {

package/dist/{backend-DaqL-MNz.js → backend-nEolvdLv.js}

@@ -150,6 +150,19 @@ function normalizeAxis(axis, ndim) {
 		return sorted(seen);
 	}
 }
+/** Check for an array of integers with no duplicates. */
+function checkInts(indices) {
+	if (typeof indices === "number") {
+		if (!Number.isInteger(indices)) throw new TypeError(`Expected integer index, got ${indices}`);
+	} else {
+		const seen = /* @__PURE__ */ new Set();
+		for (const i of indices) {
+			if (!Number.isInteger(i)) throw new TypeError(`Expected integer indices, got ${i}`);
+			if (seen.has(i)) throw new Error(`Duplicate index ${i} passed to function`);
+			seen.add(i);
+		}
+	}
+}
 function range(start, stop, step = 1) {
 	if (stop === void 0) {
 		stop = start;
@@ -896,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 = [];
@@ -918,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);
 				}
 			}
@@ -2305,10 +2318,10 @@ function tuneWebgpu(kernel) {
 			upcastedAxis.add(choices[0][2]);
 		} else break;
 	}
-	if (/chrome/i.test(navigator.userAgent) && dim.reduce < dim.unroll && (prod(dim.st.shape.slice(dim.unroll)) <= 4 || dim.unroll === dim.upcast && prod(dim.st.shape.slice(dim.upcast)) < 64)) {
+	if (!/Mobi|Android/i.test(navigator.userAgent) && dim.reduce < dim.unroll && (prod(dim.st.shape.slice(dim.unroll)) <= 4 || dim.unroll === dim.upcast && prod(dim.st.shape.slice(dim.upcast)) < 64)) {
 		const s = dim.st.shape[dim.unroll - 1];
 		if (0 < s && s <= 32) dim.applyUnroll(dim.reduce, s);
-		else for (const splits of [4]) if (s % splits === 0) {
+		else for (const splits of [8, 4]) if (s % splits === 0) {
 			dim.applyUnroll(dim.unroll - 1, splits);
 			break;
 		}
@@ -4239,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-Dh7k9io0.js");
+		const { WebGPUBackend } = await import("./webgpu-B96vzWGE.js");
 		const importantLimits = [
 			"maxBufferSize",
 			"maxComputeInvocationsPerWorkgroup",
@@ -4277,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-RSuZKvgc.js");
+		const { WebGLBackend } = await import("./webgl-DweKSWEm.js");
 		return new WebGLBackend(gl);
 	} else throw new Error(`Backend not found: ${device}`);
 }
@@ -4313,4 +4326,4 @@ var UnsupportedRoutineError = class extends Error {
 };
 
 //#endregion
-export { AluExp, AluGroup, AluOp, AluVar, DEBUG, DType, Executable, FpHash, Kernel, PPrint, Reduction, Routine, Routines, ShapeTracker, SlotError, UnsupportedOpError, UnsupportedRoutineError, accessorAluExp, accessorGlobal, assertNonNull, byteWidth, checkAxis, deepEqual, defaultDevice, devices, dtypedArray, dtypedJsArray, findPow2, generalBroadcast, getBackend, init, invertPermutation, isFloatDtype, isNumberPair, isPermutation, mapSetUnion, normalizeAxis, partitionList, prod, promoteTypes, range, recursiveFlatten, rep, runWithCache, setDebug, strip1, toposort, tuneNullopt, tuneWebgpu, unravelAlu, unzip2, zip, zipn };
+export { AluExp, AluGroup, AluOp, AluVar, DEBUG, DType, Executable, FpHash, Kernel, PPrint, Reduction, Routine, Routines, ShapeTracker, SlotError, UnsupportedOpError, UnsupportedRoutineError, accessorAluExp, accessorGlobal, assertNonNull, byteWidth, checkAxis, checkInts, deepEqual, defaultDevice, devices, dtypedArray, dtypedJsArray, findPow2, generalBroadcast, getBackend, init, invertPermutation, isFloatDtype, isNumberPair, isPermutation, mapSetUnion, normalizeAxis, partitionList, prod, promoteTypes, range, recursiveFlatten, rep, runWithCache, setDebug, strip1, toposort, tuneNullopt, tuneWebgpu, unravelAlu, unzip2, zip, zipn };