@jax-js/jax 0.1.4 → 0.1.6

package/README.md

@@ -58,13 +58,13 @@ import { numpy as np } from "@jax-js/jax";
 const ar = np.array([1, 2, 3]);
 ```
 
-By default, this is a float32 array, but you can also specify a dtype explicitly:
+By default, this is a float32 array, but you can specify a different dtype:
 
 ```ts
-const ar = np.array([1, 2, 3], { dtype: np.float32 });
+const ar = np.array([1, 2, 3], { dtype: np.int32 });
 ```
 
-For more efficient construction, you can create an array from a JS `TypedArray` buffer:
+For more efficient construction, create an array from a JS `TypedArray` buffer:
 
 ```ts
 const buf = new Float32Array([10, 20, 30, 100, 200, 300]);
@@ -223,14 +223,18 @@ Note that you need to use `type` alias syntax rather than `interface` to define
 Similar to JAX, jax-js has a concept of "devices" which are a backend that stores Arrays in memory
 and determines how to execute compiled operations on them.
 
-There are currently 3 devices in jax-js:
+There are currently 4 devices in jax-js:
 
-- `cpu`: Slow, mostly for debugging purposes.
+- `cpu`: Slow, interpreted JS, only meant for debugging.
 - `wasm`: [WebAssembly](https://webassembly.org/), currently single-threaded and blocking.
 - `webgpu`: [WebGPU](https://developer.mozilla.org/en-US/docs/Web/API/WebGPU_API), available on
   [supported browsers](https://caniuse.com/webgpu) (Chrome, Firefox, Safari, iOS).
+- `webgl`: [WebGL2](https://developer.mozilla.org/en-US/docs/Web/API/WebGL2RenderingContext), via
+  fragment shaders. This is an older graphics API that runs on almost all browsers, but it is much
+  slower than WebGPU. It's offered on a best-effort basis and not as well-supported.
 
-The default device is `wasm`, but you can change this at startup time:
+**We recommend `webgpu` for best performance, especially when running neural networks.** The default
+device is `wasm`, but you can change this at startup time:
 
 ```ts
 import { defaultDevice, init } from "@jax-js/jax";
@@ -333,7 +337,6 @@ Contributions are welcomed! Some fruitful areas to look into:
   able to generate `traceEvents` from backends (especially on GPU, with precise timestamp queries)
   to help with model performance debugging.
 - Helping the JIT compiler to fuse operations in more cases, like `tanh` branches.
-- Adding WebGL runtime for older browsers that don't support WebGPU.
 - Making a fast transformer inference engine, comparing against onnxruntime-web.
 
 You may join our [Discord server](https://discord.gg/BW6YsCd4Tf) and chat with the community.

package/dist/{backend-Bu9GY6sK.cjs → backend-D7s-Retx.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;
@@ -1471,10 +1484,37 @@ let Routines = /* @__PURE__ */ function(Routines$1) {
 	Routines$1["Sort"] = "Sort";
 	/** Returns `int32` indices of the stably sorted array. */
 	Routines$1["Argsort"] = "Argsort";
-	/** Solve a triangular system of questions. */
+	/**
+	* Solve a triangular system of equations.
+	*
+	* The first batch of inputs `A` should be of shape `[..., N, N]` and upper
+	* triangular, while the second batch `B` should be of shape `[..., M, N]`.
+	*
+	* Solves for `X` in the equation `A @ X.T = B.T`, where `A` is the
+	* triangular matrix. This is equivalent to `X = B @ A^-T`.
+	*/
 	Routines$1["TriangularSolve"] = "TriangularSolve";
-	/** Cholesky decomposition of 2D positive semi-definite matrices. */
+	/**
+	* Cholesky decomposition of 2D positive semi-definite matrices.
+	*
+	* The input batch should be of shape `[..., N, N]`, and the output batch is
+	* of the same shape, containing the lower-triangular matrix `L` such that
+	* `A = L @ L.T`. Behavior is unspecified if A is not positive semi-definite.
+	*/
 	Routines$1["Cholesky"] = "Cholesky";
+	/**
+	* LU decomposition of 2D rectangular matrices.
+	*
+	* The input is a batch of shape `[..., M, N]`, and the output is a tuple of
+	* three arrays: `LU, Pivots, Permutation`.
+	*
+	* - `LU` is of shape `[..., M, N]`, containing the combined lower and upper
+	*   triangular matrices. (lower triangular = implicit unit diagonal)
+	* - `Pivots` is of shape `[..., min(M, N)]`, containing the row swaps.
+	* - `Permutation` is of shape `[..., M]`, containing the permutation vector
+	*   such that `P = eye(M).slice(Permutation)` -> `P @ A = L @ U`.
+	*/
+	Routines$1["LU"] = "LU";
 	return Routines$1;
 }({});
 function runCpuRoutine(routine, inputs, outputs) {
@@ -1486,6 +1526,7 @@ function runCpuRoutine(routine, inputs, outputs) {
 		case Routines.Argsort: return runArgsort(type, inputAr, outputAr);
 		case Routines.TriangularSolve: return runTriangularSolve(type, inputAr, outputAr, routine.params);
 		case Routines.Cholesky: return runCholesky(type, inputAr, outputAr);
+		case Routines.LU: return runLU(type, inputAr, outputAr);
 		default:
 	}
 }
@@ -1546,6 +1587,50 @@ function runCholesky(type, [x], [y]) {
 		}
 	}
 }
+function runLU(type, [a], [lu, pivots, perm]) {
+	const shape = type.inputShapes[0];
+	if (shape.length < 2) throw new Error("lu: input must be at least 2D");
+	const m = shape[shape.length - 2];
+	const n = shape[shape.length - 1];
+	const r = Math.min(m, n);
+	for (let offset = 0; offset < a.length; offset += m * n) {
+		const ar = a.subarray(offset, offset + m * n);
+		const out = lu.subarray(offset, offset + m * n);
+		const batchIdx = offset / (m * n);
+		const piv = pivots.subarray(batchIdx * r, (batchIdx + 1) * r);
+		const p = perm.subarray(batchIdx * m, (batchIdx + 1) * m);
+		out.set(ar);
+		for (let i = 0; i < m; i++) p[i] = i;
+		for (let j = 0; j < r; j++) {
+			let maxVal = Math.abs(out[j * n + j]);
+			let maxRow = j;
+			for (let i = j + 1; i < m; i++) {
+				const val = Math.abs(out[i * n + j]);
+				if (val > maxVal) {
+					maxVal = val;
+					maxRow = i;
+				}
+			}
+			piv[j] = maxRow;
+			if (maxRow !== j) {
+				for (let col = 0; col < n; col++) {
+					const tmp = out[j * n + col];
+					out[j * n + col] = out[maxRow * n + col];
+					out[maxRow * n + col] = tmp;
+				}
+				const tmpP = p[j];
+				p[j] = p[maxRow];
+				p[maxRow] = tmpP;
+			}
+			const diag = out[j * n + j];
+			if (diag !== 0) for (let i = j + 1; i < m; i++) {
+				const factor = out[i * n + j] / diag;
+				out[i * n + j] = factor;
+				for (let col = j + 1; col < n; col++) out[i * n + col] -= factor * out[j * n + col];
+			}
+		}
+	}
+}
 
 //#endregion
 //#region src/shape.ts
@@ -2234,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;
 		}
@@ -3513,7 +3598,7 @@ var I32 = class {
 	shr_u = BINARY_OP("shr_u", 118, "i32", "i32", "i32");
 	rotl = BINARY_OP("rotl", 119, "i32", "i32", "i32");
 	rotr = BINARY_OP("rotr", 120, "i32", "i32", "i32");
-	eqz = BINARY_OP("eqz", 69, "i32", "i32", "i32");
+	eqz = UNARY_OP("eqz", 69, "i32", "i32");
 	eq = BINARY_OP("eq", 70, "i32", "i32", "i32");
 	ne = BINARY_OP("ne", 71, "i32", "i32", "i32");
 	trunc_f32_s = UNARY_OP("trunc_f32_s", 168, "f32", "i32");
@@ -3978,7 +4063,7 @@ function translateExp(cg, funcs, exp, ctx) {
 			else throw new UnsupportedOpError(op, dtype, "wasm");
 			else if (op === AluOp.Min || op === AluOp.Max) if (isFloatDtype(dtype)) if (op === AluOp.Min) dtyF(cg, op, dtype).min();
 			else dtyF(cg, op, dtype).max();
-			else if (dtype === DType.Int32 || dtype === DType.Uint32) {
+			else if (dtype === DType.Int32 || dtype === DType.Uint32 || dtype === DType.Bool) {
 				const a = cg.local.declare(cg.i32);
 				const b = cg.local.declare(cg.i32);
 				cg.local.set(b);
@@ -4128,7 +4213,8 @@ function dtyF(cg, op, dtype) {
 const devices = [
 	"cpu",
 	"wasm",
-	"webgpu"
+	"webgpu",
+	"webgl"
 ];
 const initializedBackends = /* @__PURE__ */ new Map();
 initializedBackends.set("cpu", new CpuBackend());
@@ -4167,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-Oj3Kd-kd.cjs"));
+		const { WebGPUBackend } = await Promise.resolve().then(() => require("./webgpu-rraa6dfz.cjs"));
 		const importantLimits = [
 			"maxBufferSize",
 			"maxComputeInvocationsPerWorkgroup",
@@ -4191,6 +4277,22 @@ async function createBackend(device) {
 			console.error("Unexpected error requesting WebGPU device:", error);
 			return null;
 		}
+	} else if (device === "webgl") {
+		if (typeof WebGL2RenderingContext === "undefined") return null;
+		const canvas = new OffscreenCanvas(0, 0);
+		const gl = canvas.getContext("webgl2", {
+			alpha: false,
+			antialias: false,
+			premultipliedAlpha: false,
+			preserveDrawingBuffer: false,
+			depth: false,
+			stencil: false,
+			failIfMajorPerformanceCaveat: true
+		});
+		if (!gl) return null;
+		if (!gl.getExtension("EXT_color_buffer_float")) return null;
+		const { WebGLBackend } = await Promise.resolve().then(() => require("./webgl-CyfzNW8T.cjs"));
+		return new WebGLBackend(gl);
 	} else throw new Error(`Backend not found: ${device}`);
 }
 /** Retrieve a backend that has been initialized. */
@@ -4357,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 () {
@@ -4507,6 +4615,12 @@ Object.defineProperty(exports, 'toposort', {
     return toposort;
   }
 });
+Object.defineProperty(exports, 'tuneNullopt', {
+  enumerable: true,
+  get: function () {
+    return tuneNullopt;
+  }
+});
 Object.defineProperty(exports, 'tuneWebgpu', {
   enumerable: true,
   get: function () {

package/dist/{backend-tngXtWe4.js → backend-Dx6Ob2D1.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;
@@ -1470,10 +1483,37 @@ let Routines = /* @__PURE__ */ function(Routines$1) {
 	Routines$1["Sort"] = "Sort";
 	/** Returns `int32` indices of the stably sorted array. */
 	Routines$1["Argsort"] = "Argsort";
-	/** Solve a triangular system of questions. */
+	/**
+	* Solve a triangular system of equations.
+	*
+	* The first batch of inputs `A` should be of shape `[..., N, N]` and upper
+	* triangular, while the second batch `B` should be of shape `[..., M, N]`.
+	*
+	* Solves for `X` in the equation `A @ X.T = B.T`, where `A` is the
+	* triangular matrix. This is equivalent to `X = B @ A^-T`.
+	*/
 	Routines$1["TriangularSolve"] = "TriangularSolve";
-	/** Cholesky decomposition of 2D positive semi-definite matrices. */
+	/**
+	* Cholesky decomposition of 2D positive semi-definite matrices.
+	*
+	* The input batch should be of shape `[..., N, N]`, and the output batch is
+	* of the same shape, containing the lower-triangular matrix `L` such that
+	* `A = L @ L.T`. Behavior is unspecified if A is not positive semi-definite.
+	*/
 	Routines$1["Cholesky"] = "Cholesky";
+	/**
+	* LU decomposition of 2D rectangular matrices.
+	*
+	* The input is a batch of shape `[..., M, N]`, and the output is a tuple of
+	* three arrays: `LU, Pivots, Permutation`.
+	*
+	* - `LU` is of shape `[..., M, N]`, containing the combined lower and upper
+	*   triangular matrices. (lower triangular = implicit unit diagonal)
+	* - `Pivots` is of shape `[..., min(M, N)]`, containing the row swaps.
+	* - `Permutation` is of shape `[..., M]`, containing the permutation vector
+	*   such that `P = eye(M).slice(Permutation)` -> `P @ A = L @ U`.
+	*/
+	Routines$1["LU"] = "LU";
 	return Routines$1;
 }({});
 function runCpuRoutine(routine, inputs, outputs) {
@@ -1485,6 +1525,7 @@ function runCpuRoutine(routine, inputs, outputs) {
 		case Routines.Argsort: return runArgsort(type, inputAr, outputAr);
 		case Routines.TriangularSolve: return runTriangularSolve(type, inputAr, outputAr, routine.params);
 		case Routines.Cholesky: return runCholesky(type, inputAr, outputAr);
+		case Routines.LU: return runLU(type, inputAr, outputAr);
 		default:
 	}
 }
@@ -1545,6 +1586,50 @@ function runCholesky(type, [x], [y]) {
 		}
 	}
 }
+function runLU(type, [a], [lu, pivots, perm]) {
+	const shape = type.inputShapes[0];
+	if (shape.length < 2) throw new Error("lu: input must be at least 2D");
+	const m = shape[shape.length - 2];
+	const n = shape[shape.length - 1];
+	const r = Math.min(m, n);
+	for (let offset = 0; offset < a.length; offset += m * n) {
+		const ar = a.subarray(offset, offset + m * n);
+		const out = lu.subarray(offset, offset + m * n);
+		const batchIdx = offset / (m * n);
+		const piv = pivots.subarray(batchIdx * r, (batchIdx + 1) * r);
+		const p = perm.subarray(batchIdx * m, (batchIdx + 1) * m);
+		out.set(ar);
+		for (let i = 0; i < m; i++) p[i] = i;
+		for (let j = 0; j < r; j++) {
+			let maxVal = Math.abs(out[j * n + j]);
+			let maxRow = j;
+			for (let i = j + 1; i < m; i++) {
+				const val = Math.abs(out[i * n + j]);
+				if (val > maxVal) {
+					maxVal = val;
+					maxRow = i;
+				}
+			}
+			piv[j] = maxRow;
+			if (maxRow !== j) {
+				for (let col = 0; col < n; col++) {
+					const tmp = out[j * n + col];
+					out[j * n + col] = out[maxRow * n + col];
+					out[maxRow * n + col] = tmp;
+				}
+				const tmpP = p[j];
+				p[j] = p[maxRow];
+				p[maxRow] = tmpP;
+			}
+			const diag = out[j * n + j];
+			if (diag !== 0) for (let i = j + 1; i < m; i++) {
+				const factor = out[i * n + j] / diag;
+				out[i * n + j] = factor;
+				for (let col = j + 1; col < n; col++) out[i * n + col] -= factor * out[j * n + col];
+			}
+		}
+	}
+}
 
 //#endregion
 //#region src/shape.ts
@@ -2233,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;
 		}
@@ -3512,7 +3597,7 @@ var I32 = class {
 	shr_u = BINARY_OP("shr_u", 118, "i32", "i32", "i32");
 	rotl = BINARY_OP("rotl", 119, "i32", "i32", "i32");
 	rotr = BINARY_OP("rotr", 120, "i32", "i32", "i32");
-	eqz = BINARY_OP("eqz", 69, "i32", "i32", "i32");
+	eqz = UNARY_OP("eqz", 69, "i32", "i32");
 	eq = BINARY_OP("eq", 70, "i32", "i32", "i32");
 	ne = BINARY_OP("ne", 71, "i32", "i32", "i32");
 	trunc_f32_s = UNARY_OP("trunc_f32_s", 168, "f32", "i32");
@@ -3977,7 +4062,7 @@ function translateExp(cg, funcs, exp, ctx) {
 			else throw new UnsupportedOpError(op, dtype, "wasm");
 			else if (op === AluOp.Min || op === AluOp.Max) if (isFloatDtype(dtype)) if (op === AluOp.Min) dtyF(cg, op, dtype).min();
 			else dtyF(cg, op, dtype).max();
-			else if (dtype === DType.Int32 || dtype === DType.Uint32) {
+			else if (dtype === DType.Int32 || dtype === DType.Uint32 || dtype === DType.Bool) {
 				const a = cg.local.declare(cg.i32);
 				const b = cg.local.declare(cg.i32);
 				cg.local.set(b);
@@ -4127,7 +4212,8 @@ function dtyF(cg, op, dtype) {
 const devices = [
 	"cpu",
 	"wasm",
-	"webgpu"
+	"webgpu",
+	"webgl"
 ];
 const initializedBackends = /* @__PURE__ */ new Map();
 initializedBackends.set("cpu", new CpuBackend());
@@ -4166,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-ChVgx3b6.js");
+		const { WebGPUBackend } = await import("./webgpu-C-VfevQW.js");
 		const importantLimits = [
 			"maxBufferSize",
 			"maxComputeInvocationsPerWorkgroup",
@@ -4190,6 +4276,22 @@ async function createBackend(device) {
 			console.error("Unexpected error requesting WebGPU device:", error);
 			return null;
 		}
+	} else if (device === "webgl") {
+		if (typeof WebGL2RenderingContext === "undefined") return null;
+		const canvas = new OffscreenCanvas(0, 0);
+		const gl = canvas.getContext("webgl2", {
+			alpha: false,
+			antialias: false,
+			premultipliedAlpha: false,
+			preserveDrawingBuffer: false,
+			depth: false,
+			stencil: false,
+			failIfMajorPerformanceCaveat: true
+		});
+		if (!gl) return null;
+		if (!gl.getExtension("EXT_color_buffer_float")) return null;
+		const { WebGLBackend } = await import("./webgl-CLLvzJlO.js");
+		return new WebGLBackend(gl);
 	} else throw new Error(`Backend not found: ${device}`);
 }
 /** Retrieve a backend that has been initialized. */
@@ -4224,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, 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 };