@jax-js/jax 0.1.3 → 0.1.5

package/README.md

@@ -3,7 +3,8 @@
 <p align="center"><strong>
   <a href="https://jax-js.com">Website</a> |
   <a href="https://jax-js.com/docs/">API Reference</a> |
-  <a href="./FEATURES.md">Compatibility Table</a>
+  <a href="./FEATURES.md">Compatibility Table</a> |
+  <a href="https://discord.gg/BW6YsCd4Tf">Discord</a>
 </strong></p>
 
 **jax-js** is a machine learning framework for the browser. It aims to bring
@@ -57,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]);
@@ -222,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";
@@ -323,7 +328,7 @@ pnpm -C website dev
 
 ## Future work / help wanted
 
-Contributions are welcomed! Especially in:
+Contributions are welcomed! Some fruitful areas to look into:
 
 - Adding support for more JAX functions and operations, see [compatibility table](./FEATURES.md).
 - Improving performance of the WebGPU and Wasm runtimes, generating better kernels, and using SIMD
@@ -332,5 +337,6 @@ Contributions are welcomed! Especially in:
   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-BY8wlLEl.js → backend-DaqL-MNz.js}

@@ -68,6 +68,9 @@ function zipn(...arrays) {
 	const minLength = Math.min(...arrays.map((x) => x.length));
 	return Array.from({ length: minLength }, (_, i) => arrays.map((arr) => arr[i]));
 }
+function sorted(arr) {
+	return [...arr].sort((a, b) => a - b);
+}
 function rep(length, value) {
 	if (value instanceof Function) return new Array(length).fill(0).map((_, i) => value(i));
 	return new Array(length).fill(value);
@@ -144,7 +147,7 @@ function normalizeAxis(axis, ndim) {
 			if (seen.has(ca)) throw new Error(`Duplicate axis ${ca} passed to function`);
 			seen.add(ca);
 		}
-		return [...seen].sort();
+		return sorted(seen);
 	}
 }
 function range(start, stop, step = 1) {
@@ -1327,7 +1330,7 @@ var Reduction = class {
 	/** Evaluate this operation on CPU. */
 	evaluate(...values) {
 		if (this.dtype === DType.Bool) {
-			if (this.op === AluOp.Add || this.op === AluOp.Max) return values.reduce((a, b) => a || b, true);
+			if (this.op === AluOp.Add || this.op === AluOp.Max) return values.reduce((a, b) => a || b, false);
 			else if (this.op === AluOp.Mul || this.op === AluOp.Min) return values.reduce((a, b) => a && b, true);
 		} else if (this.dtype === DType.Int32) {
 			if (this.op === AluOp.Add) return values.reduce((a, b) => a + b | 0, 0);
@@ -1437,6 +1440,184 @@ function erfc(x) {
 	else return 2 - _erfapprox$1(-x);
 }
 
+//#endregion
+//#region src/routine.ts
+/**
+* Advanced operations that don't fit into the `AluExp` compiler representation.
+*
+* Some routines like iterative matrix algorithms, FFTs, or sorting may not be
+* easy to express efficiently as a `Kernel` object. These also tend to be
+* somewhat expensive, so the benefit of kernel fusion and inlining is less
+* relevant.
+*
+* For these operations, we dispatch them as a custom operation on the backend,
+* which each backend implements in a specific way. These are listed in the
+* `Routines` enum below.
+*
+* Routines cannot be fused into other kernels and always operate on contiguous
+* arrays (default `ShapeTracker`).
+*/
+var Routine = class {
+	constructor(name, type, params) {
+		this.name = name;
+		this.type = type;
+		this.params = params;
+	}
+};
+/** One of the valid `Routine` that can be dispatched to backend. */
+let Routines = /* @__PURE__ */ function(Routines$1) {
+	/** Stable sorting algorithm along the last axis. */
+	Routines$1["Sort"] = "Sort";
+	/** Returns `int32` indices of the stably sorted array. */
+	Routines$1["Argsort"] = "Argsort";
+	/**
+	* 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.
+	*
+	* 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) {
+	const { name, type } = routine;
+	const inputAr = inputs.map((buf, i) => dtypedArray(type.inputDtypes[i], buf));
+	const outputAr = outputs.map((buf, i) => dtypedArray(type.outputDtypes[i], buf));
+	switch (name) {
+		case Routines.Sort: return runSort(type, inputAr, outputAr);
+		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:
+	}
+}
+function runSort(type, [x], [y]) {
+	const xs = type.inputShapes[0];
+	if (xs.length === 0) throw new Error("sort: cannot sort a scalar");
+	const n = xs[xs.length - 1];
+	y.set(x);
+	for (let i = 0; i < y.length; i += n) y.subarray(i, i + n).sort();
+}
+function runArgsort(type, [x], [y, yi]) {
+	const xs = type.inputShapes[0];
+	if (xs.length === 0) throw new Error("argsort: cannot sort a scalar");
+	const n = xs[xs.length - 1];
+	for (let offset = 0; offset < y.length; offset += n) {
+		const ar = x.subarray(offset, offset + n);
+		const out = y.subarray(offset, offset + n);
+		const outi = yi.subarray(offset, offset + n);
+		for (let i = 0; i < n; i++) outi[i] = i;
+		outi.sort((a, b) => ar[a] - ar[b]);
+		for (let i = 0; i < n; i++) out[i] = ar[outi[i]];
+	}
+}
+function runTriangularSolve(type, [a, b], [x], { unitDiagonal }) {
+	const as = type.inputShapes[0];
+	const bs = type.inputShapes[1];
+	if (as.length < 2) throw new Error(`triangular_solve: a must be at least 2D, got ${as}`);
+	if (bs.length < 2) throw new Error(`triangular_solve: b must be at least 2D, got ${bs}`);
+	const n = as[as.length - 2];
+	if (n !== as[as.length - 1] || n !== bs[bs.length - 1]) throw new Error(`triangular_solve: incompatible shapes a=${as}, b=${bs}`);
+	const batch = bs[bs.length - 2];
+	for (let counter = 0; counter < a.length / (n * n); counter++) {
+		const a1 = a.subarray(counter * n * n, (counter + 1) * n * n);
+		for (let t = 0; t < batch; t++) {
+			const b1 = b.subarray((counter * batch + t) * n, (counter * batch + t + 1) * n);
+			const x1 = x.subarray((counter * batch + t) * n, (counter * batch + t + 1) * n);
+			for (let i = n - 1; i >= 0; i--) {
+				let sum = b1[i];
+				for (let j = i + 1; j < n; j++) sum -= a1[i * n + j] * x1[j];
+				x1[i] = unitDiagonal ? sum : sum / a1[i * n + i];
+			}
+		}
+	}
+}
+function runCholesky(type, [x], [y]) {
+	const xs = type.inputShapes[0];
+	if (xs.length < 2) throw new Error("cholesky: input must be at least 2D");
+	const n = xs[xs.length - 2];
+	const m = xs[xs.length - 1];
+	if (n !== m) throw new Error(`cholesky: input must be square, got [${n}, ${m}]`);
+	for (let offset = 0; offset < y.length; offset += n * n) {
+		const ar = x.subarray(offset, offset + n * n);
+		const out = y.subarray(offset, offset + n * n);
+		for (let i = 0; i < n; i++) for (let j = 0; j <= i; j++) {
+			let sum = ar[i * n + j];
+			for (let k = 0; k < j; k++) sum -= out[i * n + k] * out[j * n + k];
+			out[i * n + j] = i === j ? Math.sqrt(sum) : sum / out[j * n + j];
+		}
+	}
+}
+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
 const jstr = JSON.stringify;
@@ -1907,7 +2088,7 @@ var ShapeTracker = class ShapeTracker {
 		let st = this;
 		if (axis.length > 0) {
 			const unsqueezed = [...st.shape];
-			for (const i of axis.toSorted()) unsqueezed.splice(i, 0, 1);
+			for (const i of sorted(axis)) unsqueezed.splice(i, 0, 1);
 			st = st.reshape(unsqueezed);
 		}
 		return st.expand(newShape);
@@ -2132,7 +2313,7 @@ function tuneWebgpu(kernel) {
 			break;
 		}
 	}
-	for (const ax of Array.from(upcastedAxis).sort()) {
+	for (const ax of sorted(upcastedAxis)) {
 		const s = dim.st.shape[ax];
 		for (const amount of [8, 4]) if (s % amount === 0) {
 			dim.applyLocal(ax, amount);
@@ -2250,13 +2431,21 @@ var CpuBackend = class {
 		if (count === void 0) count = buffer.byteLength - start;
 		return buffer.slice(start, start + count);
 	}
-	async prepare(kernel) {
-		return this.prepareSync(kernel);
+	async prepareKernel(kernel) {
+		return this.prepareKernelSync(kernel);
 	}
-	prepareSync(kernel) {
+	prepareKernelSync(kernel) {
 		return new Executable(kernel, void 0);
 	}
-	dispatch({ kernel }, inputs, outputs) {
+	async prepareRoutine(routine) {
+		return this.prepareRoutineSync(routine);
+	}
+	prepareRoutineSync(routine) {
+		return new Executable(routine, void 0);
+	}
+	dispatch(exe, inputs, outputs) {
+		if (exe.source instanceof Routine) return runCpuRoutine(exe.source, inputs.map((slot) => this.#getBuffer(slot)), outputs.map((slot) => this.#getBuffer(slot)));
+		const kernel = exe.source;
 		const { exp, epilogue } = tuneNullopt(kernel);
 		const inputBuffers = inputs.map((slot) => this.#getBuffer(slot));
 		const outputBuffers = outputs.map((slot) => this.#getBuffer(slot));
@@ -2314,8 +2503,10 @@ var WasmAllocator = class {
 		const sizeClass = this.#findSizeClass(size);
 		const freeList = this.#freeLists.get(sizeClass);
 		let ptr;
-		if (freeList && freeList.length > 0) ptr = freeList.pop();
-		else ptr = this.#bumpAlloc(sizeClass);
+		if (freeList && freeList.length > 0) {
+			ptr = freeList.pop();
+			new Uint8Array(this.#memory.buffer, ptr, sizeClass).fill(0);
+		} else ptr = this.#bumpAlloc(sizeClass);
 		this.#allocatedBuffers.set(ptr, sizeClass);
 		return ptr;
 	}
@@ -3393,7 +3584,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");
@@ -3681,10 +3872,10 @@ var WasmBackend = class {
 		if (count === void 0) count = buffer.byteLength - start;
 		return buffer.slice(start, start + count);
 	}
-	async prepare(kernel) {
-		return this.prepareSync(kernel);
+	async prepareKernel(kernel) {
+		return this.prepareKernelSync(kernel);
 	}
-	prepareSync(kernel) {
+	prepareKernelSync(kernel) {
 		const kernelHash = FpHash.hash(kernel);
 		const module = runWithCache(moduleCache, kernelHash.toString(), () => {
 			const bytes = codegenWasm(kernel);
@@ -3692,7 +3883,14 @@ var WasmBackend = class {
 		});
 		return new Executable(kernel, { module });
 	}
+	async prepareRoutine(routine) {
+		return this.prepareRoutineSync(routine);
+	}
+	prepareRoutineSync(routine) {
+		return new Executable(routine, void 0);
+	}
 	dispatch(exe, inputs, outputs) {
+		if (exe.source instanceof Routine) return runCpuRoutine(exe.source, inputs.map((slot) => this.#getBuffer(slot)), outputs.map((slot) => this.#getBuffer(slot)));
 		const instance = new WebAssembly.Instance(exe.data.module, { env: { memory: this.#memory } });
 		const func = instance.exports.kernel;
 		const ptrs = [...inputs, ...outputs].map((slot) => this.#buffers.get(slot).ptr);
@@ -3851,7 +4049,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);
@@ -4001,7 +4199,8 @@ function dtyF(cg, op, dtype) {
 const devices = [
 	"cpu",
 	"wasm",
-	"webgpu"
+	"webgpu",
+	"webgl"
 ];
 const initializedBackends = /* @__PURE__ */ new Map();
 initializedBackends.set("cpu", new CpuBackend());
@@ -4040,7 +4239,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-C9iAP5h5.js");
+		const { WebGPUBackend } = await import("./webgpu-Dh7k9io0.js");
 		const importantLimits = [
 			"maxBufferSize",
 			"maxComputeInvocationsPerWorkgroup",
@@ -4064,6 +4263,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-RSuZKvgc.js");
+		return new WebGLBackend(gl);
 	} else throw new Error(`Backend not found: ${device}`);
 }
 /** Retrieve a backend that has been initialized. */
@@ -4074,8 +4289,8 @@ function getBackend(device) {
 	return backend;
 }
 var Executable = class {
-	constructor(kernel, data) {
-		this.kernel = kernel;
+	constructor(source, data) {
+		this.source = source;
 		this.data = data;
 	}
 };
@@ -4091,7 +4306,11 @@ var UnsupportedOpError = class extends Error {
 		super(msg);
 	}
 };
+var UnsupportedRoutineError = class extends Error {
+	constructor(name, device) {
+		super(`routine '${name}' is not supported in ${device} backend`);
+	}
+};
 
 //#endregion
-export { AluExp, AluGroup, AluOp, AluVar, DEBUG, DType, Executable, FpHash, Kernel, PPrint, Reduction, ShapeTracker, SlotError, UnsupportedOpError, 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 };
-//# sourceMappingURL=backend-BY8wlLEl.js.map
+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 };