@jax-js/jax 0.0.4 → 0.1.0

package/dist/{backend-EBRGmEYw.js → backend-DwIAd0AG.js}

@@ -51,6 +51,7 @@ let DEBUG = 0;
 function setDebug(level) {
 	DEBUG = level;
 }
+function assertNonNull(value) {}
 function unzip2(pairs) {
 	const lst1 = [];
 	const lst2 = [];
@@ -96,10 +97,15 @@ function deepEqual(a, b) {
 	for (const key of Object.keys(a)) if (!deepEqual(a[key], b[key])) return false;
 	return true;
 }
-function union(...sets) {
-	const result = /* @__PURE__ */ new Set();
-	for (const s of sets) if (s) for (const x of s) result.add(x);
-	return result;
+/** Produces a union of maps of sets. This mutates `a`. */
+function mapSetUnion(a, b) {
+	if (!b) return a;
+	for (const [key, setB] of b.entries()) {
+		const setA = a.get(key);
+		if (setA) for (const val of setB) setA.add(val);
+		else a.set(key, setB);
+	}
+	return a;
 }
 /** Splits the list based on a condition, `false` first then `true`. */
 function partitionList(which, array) {
@@ -206,6 +212,35 @@ function findPow2(hint, max) {
 	while (ret < hint && 2 * ret <= max) ret *= 2;
 	return ret;
 }
+/**
+* Implements a NumPy-style generalized broadcast rule on two array shapes.
+*
+* "When operating on two arrays, NumPy compares their shapes element-wise. It
+* starts with the trailing (i.e. rightmost) dimension and works its way left.
+* Two dimensions are compatible when:
+*   1. they are equal, or
+*   2. one of them is 1."
+*
+* Throws a TypeError if the broadcast is not possible.
+*
+* <https://numpy.org/doc/stable/user/basics.broadcasting.html#general-broadcasting-rules>
+*/
+function generalBroadcast(a, b) {
+	const out = [];
+	let i = a.length - 1;
+	let j = b.length - 1;
+	for (; i >= 0 && j >= 0; i--, j--) {
+		const x = a[i];
+		const y = b[j];
+		if (x === y) out.push(x);
+		else if (x === 1) out.push(y);
+		else if (y === 1) out.push(x);
+		else throw new TypeError(`Incompatible array broadcast shapes: ${a} vs ${b}`);
+	}
+	for (; i >= 0; i--) out.push(a[i]);
+	for (; j >= 0; j--) out.push(b[j]);
+	return out.reverse();
+}
 function recursiveFlatten(ar) {
 	if (!Array.isArray(ar)) return [ar];
 	return ar.flat(Infinity);
@@ -294,12 +329,12 @@ const isFloatDtype = (dtype) => dtype === DType.Float32 || dtype === DType.Float
 * **Type lattice:**
 * ```text
 * bool -> uint32 -> int32 -> float16 -> float32
-*  weak f* --^
+*  weakType --^
 * ```
 *
-* The asterisk f* is a weak type used for JS number constants. When creating
-* arrays, JS numbers default to float32 but "weak" so they cast to the dtype of
-* any array they are first combined with.
+* `weakType` represents weakly typed arrays. These are created for JS numbers,
+* which default to float32 but "weak" so they cast to the dtype of any array
+* they are first combined with, except `bool`.
 *
 * **Examples:**
 * - `promoteTypes(bool, int32) → int32`
@@ -400,6 +435,12 @@ var AluExp = class AluExp {
 	static log(a) {
 		return new AluExp(AluOp.Log, a.dtype, [a]);
 	}
+	static erf(a) {
+		return new AluExp(AluOp.Erf, a.dtype, [a]);
+	}
+	static erfc(a) {
+		return new AluExp(AluOp.Erfc, a.dtype, [a]);
+	}
 	static sqrt(a) {
 		return new AluExp(AluOp.Sqrt, a.dtype, [a]);
 	}
@@ -570,6 +611,12 @@ var AluExp = class AluExp {
 			case AluOp.Log:
 				ret = [Math.log(src[0].min), Math.log(src[0].max)];
 				break;
+			case AluOp.Erf:
+				ret = [erf(src[0].min), erf(src[0].max)];
+				break;
+			case AluOp.Erfc:
+				ret = [erfc(src[0].max), erfc(src[0].min)];
+				break;
 			case AluOp.Sqrt:
 				ret = [Math.sqrt(src[0].min), Math.sqrt(src[0].max)];
 				break;
@@ -891,6 +938,8 @@ var AluExp = class AluExp {
 				case AluOp.Atan: return Math.atan(x);
 				case AluOp.Exp: return Math.exp(x);
 				case AluOp.Log: return Math.log(x);
+				case AluOp.Erf: return erf(x);
+				case AluOp.Erfc: return erfc(x);
 				case AluOp.Sqrt: return Math.sqrt(x);
 				case AluOp.Reciprocal: return 1 / x;
 				case AluOp.Cast: {
@@ -1039,11 +1088,15 @@ var AluExp = class AluExp {
 		});
 		return result;
 	}
-	/** Produce a list of all distinct AluOp in this expression. */
+	/** Produce all distinct AluOp in this expression, with their dtypes. */
 	distinctOps() {
-		const ops = /* @__PURE__ */ new Set();
+		const ops = /* @__PURE__ */ new Map();
 		this.fold((exp) => {
-			ops.add(exp.op);
+			const s = ops.get(exp.op) ?? /* @__PURE__ */ new Set();
+			if (!s.has(exp.dtype)) {
+				s.add(exp.dtype);
+				ops.set(exp.op, s);
+			}
 		});
 		return ops;
 	}
@@ -1072,6 +1125,8 @@ let AluOp = /* @__PURE__ */ function(AluOp$1) {
 	AluOp$1["Atan"] = "Atan";
 	AluOp$1["Exp"] = "Exp";
 	AluOp$1["Log"] = "Log";
+	AluOp$1["Erf"] = "Erf";
+	AluOp$1["Erfc"] = "Erfc";
 	AluOp$1["Sqrt"] = "Sqrt";
 	AluOp$1["Reciprocal"] = "Reciprocal";
 	AluOp$1["Cast"] = "Cast";
@@ -1104,6 +1159,8 @@ const AluGroup = {
 		AluOp.Atan,
 		AluOp.Exp,
 		AluOp.Log,
+		AluOp.Erf,
+		AluOp.Erfc,
 		AluOp.Sqrt,
 		AluOp.Reciprocal,
 		AluOp.Cast,
@@ -1129,6 +1186,8 @@ const AluGroup = {
 		AluOp.Atan,
 		AluOp.Exp,
 		AluOp.Log,
+		AluOp.Erf,
+		AluOp.Erfc,
 		AluOp.Sqrt,
 		AluOp.Reciprocal
 	])
@@ -1304,6 +1363,44 @@ function threefry2x32(k0, k1, c0, c1) {
 	x1 = x1 + ks0 + 5 >>> 0;
 	return [x0, x1];
 }
+/**
+* Abramowitz & Stegun’s widely used approximation for erf(x).
+*
+* `erf(x) = 1 - P(t) * exp(-x^2)` for `x >= 0`, where `t = 1/(1 + p*x)` and
+* `P(t) = a1*t + a2*t^2 + a3*t^3 + a4*t^4 + a5*t^5`.
+*
+* Coefficients:
+*  - p = 0.3275911
+*  - a1 = 0.254829592
+*  - a2 = -0.284496736
+*  - a3 = 1.421413741
+*  - a4 = -1.453152027
+*  - a5 = 1.061405429
+*
+* This function computes just `E = P(t) * exp(-x^2)` for numerical reasons. The
+* input is assumed to be non-negative.
+*
+* Reference: https://en.wikipedia.org/wiki/Error_function#Approximation_with_elementary_functions
+*/
+function _erfapprox$1(x) {
+	const p = .3275911;
+	const a1 = .254829592;
+	const a2 = -.284496736;
+	const a3 = 1.421413741;
+	const a4 = -1.453152027;
+	const a5 = 1.061405429;
+	const t = 1 / (1 + p * x);
+	const P_t = ((((a5 * t + a4) * t + a3) * t + a2) * t + a1) * t;
+	return P_t * Math.exp(-x * x);
+}
+function erf(x) {
+	if (x >= 0) return 1 - _erfapprox$1(x);
+	else return _erfapprox$1(-x) - 1;
+}
+function erfc(x) {
+	if (x >= 0) return _erfapprox$1(x);
+	else return 2 - _erfapprox$1(-x);
+}
 
 //#endregion
 //#region src/shape.ts
@@ -1989,7 +2086,7 @@ function tuneWebgpu(kernel) {
 	}
 	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)) {
 		const s = dim.st.shape[dim.unroll - 1];
-		if (s <= 32) dim.applyUnroll(dim.reduce, s);
+		if (0 < s && s <= 32) dim.applyUnroll(dim.reduce, s);
 		else for (const splits of [4]) if (s % splits === 0) {
 			dim.applyUnroll(dim.unroll - 1, splits);
 			break;
@@ -2208,6 +2305,19 @@ var WasmAllocator = class {
 
 //#endregion
 //#region src/backend/wasm/builtins.ts
+/** Given a local `x`, evaluate `sum[i](a_i * x^i)` and push to stack. */
+function _poly(cg, x, as) {
+	if (as.length === 0) throw new Error("_poly needs at least one coefficient");
+	cg.f32.const(as[as.length - 1]);
+	for (let i = as.length - 2; i >= 0; i--) {
+		cg.local.get(x);
+		cg.f32.mul();
+		if (as[i] !== 0) {
+			cg.f32.const(as[i]);
+			cg.f32.add();
+		}
+	}
+}
 /**
 * Approximate e^x.
 *
@@ -2248,27 +2358,15 @@ function wasm_exp(cg) {
 		cg.f32.mul();
 		cg.f32.sub();
 		cg.local.set(r);
-		cg.f32.const(1 / 120);
-		cg.local.get(r);
-		cg.f32.mul();
-		cg.f32.const(1 / 24);
-		cg.f32.add();
-		cg.local.get(r);
-		cg.f32.mul();
-		cg.f32.const(1 / 6);
-		cg.f32.add();
-		cg.local.get(r);
-		cg.f32.mul();
-		cg.f32.const(1 / 2);
-		cg.f32.add();
-		cg.local.get(r);
-		cg.f32.mul();
-		cg.f32.const(1);
-		cg.f32.add();
-		cg.local.get(r);
-		cg.f32.mul();
-		cg.f32.const(1);
-		cg.f32.add();
+		_poly(cg, r, [
+			1,
+			1,
+			1 / 2,
+			1 / 6,
+			1 / 24,
+			1 / 120,
+			1 / 720
+		]);
 		cg.local.set(p);
 		cg.local.get(k);
 		cg.i32.const(127);
@@ -2296,11 +2394,6 @@ function wasm_log(cg) {
 		const m = cg.local.declare(cg.f32);
 		const t = cg.local.declare(cg.f32);
 		const t2 = cg.local.declare(cg.f32);
-		const t3 = cg.local.declare(cg.f32);
-		const t5 = cg.local.declare(cg.f32);
-		const t7 = cg.local.declare(cg.f32);
-		const lnm = cg.local.declare(cg.f32);
-		const el2 = cg.local.declare(cg.f32);
 		cg.local.get(0);
 		cg.f32.const(0);
 		cg.f32.le();
@@ -2337,41 +2430,18 @@ function wasm_log(cg) {
 		cg.local.get(t);
 		cg.f32.mul();
 		cg.local.set(t2);
+		_poly(cg, t2, [
+			2,
+			2 / 3,
+			2 / 5,
+			2 / 7
+		]);
 		cg.local.get(t);
-		cg.local.get(t2);
-		cg.f32.mul();
-		cg.local.set(t3);
-		cg.local.get(t3);
-		cg.local.get(t2);
-		cg.f32.mul();
-		cg.local.set(t5);
-		cg.local.get(t5);
-		cg.local.get(t2);
-		cg.f32.mul();
-		cg.local.set(t7);
-		cg.local.get(t7);
-		cg.f32.const(1 / 7);
-		cg.f32.mul();
-		cg.local.get(t5);
-		cg.f32.const(1 / 5);
-		cg.f32.mul();
-		cg.f32.add();
-		cg.local.get(t3);
-		cg.f32.const(1 / 3);
-		cg.f32.mul();
-		cg.f32.add();
-		cg.local.get(t);
-		cg.f32.add();
-		cg.f32.const(2);
 		cg.f32.mul();
-		cg.local.set(lnm);
 		cg.local.get(e);
 		cg.f32.convert_i32_s();
 		cg.f32.const(Math.LN2);
 		cg.f32.mul();
-		cg.local.set(el2);
-		cg.local.get(el2);
-		cg.local.get(lnm);
 		cg.f32.add();
 	});
 }
@@ -2381,7 +2451,7 @@ function wasm_log(cg) {
 * Method: reduce to y in [-π, π], then quadrant via q = round(y/(π/2))
 *         z = y - q*(π/2); use one of two polynomials on z:
 *         sin(z) ≈ z + z^3*(-1/6) + z^5*(1/120) + z^7*(-1/5040)
-*         cos(z) ≈ 1 + z^2*(-1/2) + z^4*(1/24) + z^6*(-1/720)
+*         cos(z) ≈ 1 + z^2*(-1/2) + z^4*(1/24) + z^6*(-1/720) + z^8*(1/40320)
 */
 function _sincos(cg) {
 	const y = cg.local.declare(cg.f32);
@@ -2417,35 +2487,22 @@ function _sincos(cg) {
 	cg.local.get(z);
 	cg.f32.mul();
 	cg.local.set(z2);
-	cg.f32.const(-1 / 5040);
-	cg.local.get(z2);
-	cg.f32.mul();
-	cg.f32.const(1 / 120);
-	cg.f32.add();
-	cg.local.get(z2);
-	cg.f32.mul();
-	cg.f32.const(-1 / 6);
-	cg.f32.add();
-	cg.local.get(z2);
-	cg.f32.mul();
-	cg.f32.const(1);
-	cg.f32.add();
+	_poly(cg, z2, [
+		1,
+		-1 / 6,
+		1 / 120,
+		-1 / 5040
+	]);
 	cg.local.get(z);
 	cg.f32.mul();
 	cg.local.set(sz);
-	cg.f32.const(-1 / 720);
-	cg.local.get(z2);
-	cg.f32.mul();
-	cg.f32.const(1 / 24);
-	cg.f32.add();
-	cg.local.get(z2);
-	cg.f32.mul();
-	cg.f32.const(-1 / 2);
-	cg.f32.add();
-	cg.local.get(z2);
-	cg.f32.mul();
-	cg.f32.const(1);
-	cg.f32.add();
+	_poly(cg, z2, [
+		1,
+		-1 / 2,
+		1 / 24,
+		-1 / 720,
+		1 / 40320
+	]);
 	cg.local.set(cz);
 	return {
 		q,
@@ -2527,24 +2584,16 @@ function _atan(cg) {
 	cg.local.get(z);
 	cg.f32.mul();
 	cg.local.set(z2);
-	cg.f32.const(.0415796528637);
-	cg.local.get(z2);
-	cg.f32.mul();
-	cg.f32.const(.661705427875);
-	cg.f32.add();
-	cg.local.get(z2);
-	cg.f32.mul();
-	cg.f32.const(.999998614341);
-	cg.f32.add();
-	cg.f32.const(.173698870181);
-	cg.local.get(z2);
-	cg.f32.mul();
-	cg.f32.const(.994987933645);
-	cg.f32.add();
-	cg.local.get(z2);
-	cg.f32.mul();
-	cg.f32.const(1);
-	cg.f32.add();
+	_poly(cg, z2, [
+		.999998614341,
+		.661705427875,
+		.0415796528637
+	]);
+	_poly(cg, z2, [
+		1,
+		.994987933645,
+		.173698870181
+	]);
 	cg.f32.div();
 	cg.local.get(z);
 	cg.f32.mul();
@@ -2598,6 +2647,74 @@ function wasm_asin(cg) {
 	});
 }
 /**
+* Helper function for erf/erfc approximation.
+*
+* See `_erfapprox` in alu.ts for details on the algorithm used.
+*/
+function _erfapprox(cg, exp_func) {
+	const x = cg.local.declare(cg.f32);
+	const t = cg.local.declare(cg.f32);
+	cg.local.set(x);
+	const p = .3275911;
+	const a1 = .254829592;
+	const a2 = -.284496736;
+	const a3 = 1.421413741;
+	const a4 = -1.453152027;
+	const a5 = 1.061405429;
+	cg.f32.const(1);
+	cg.f32.const(1);
+	cg.f32.const(p);
+	cg.local.get(x);
+	cg.f32.mul();
+	cg.f32.add();
+	cg.f32.div();
+	cg.local.set(t);
+	_poly(cg, t, [
+		0,
+		a1,
+		a2,
+		a3,
+		a4,
+		a5
+	]);
+	cg.local.get(x);
+	cg.f32.neg();
+	cg.local.get(x);
+	cg.f32.mul();
+	cg.call(exp_func);
+	cg.f32.mul();
+}
+/** Approximate erf(x) (error function). */
+function wasm_erf(cg, exp) {
+	return cg.function([cg.f32], [cg.f32], () => {
+		cg.f32.const(1);
+		cg.local.get(0);
+		cg.f32.abs();
+		_erfapprox(cg, exp);
+		cg.f32.sub();
+		cg.local.get(0);
+		cg.f32.copysign();
+	});
+}
+/** Approximate erfc(x) (complementary error function). */
+function wasm_erfc(cg, exp) {
+	return cg.function([cg.f32], [cg.f32], () => {
+		const e = cg.local.declare(cg.f32);
+		cg.local.get(0);
+		cg.f32.abs();
+		_erfapprox(cg, exp);
+		cg.local.set(e);
+		cg.f32.const(2);
+		cg.local.get(e);
+		cg.f32.sub();
+		cg.local.get(e);
+		cg.local.get(0);
+		cg.f32.const(0);
+		cg.f32.lt();
+		cg.select();
+	});
+}
+/**
 * Threefry2x32 pseudorandom number generator.
 *
 * Takes two 32-bit keys and two 32-bit counters as input,
@@ -3472,14 +3589,16 @@ function codegenWasm(kernel) {
 	if (DEBUG >= 3) console.info(`kernel.exp: ${kernel.exp}\ntune.exp: ${tune.exp}`);
 	const cg = new CodeGenerator();
 	cg.memory.import("env", "memory");
-	const distinctOps = union(tune.exp.distinctOps(), re?.epilogue.distinctOps());
+	const distinctOps = mapSetUnion(tune.exp.distinctOps(), re?.epilogue.distinctOps());
 	const funcs = {};
 	if (distinctOps.has(AluOp.Sin)) funcs.sin = wasm_sin(cg);
 	if (distinctOps.has(AluOp.Cos)) funcs.cos = wasm_cos(cg);
 	if (distinctOps.has(AluOp.Asin)) funcs.asin = wasm_asin(cg);
 	if (distinctOps.has(AluOp.Atan)) funcs.atan = wasm_atan(cg);
-	if (distinctOps.has(AluOp.Exp)) funcs.exp = wasm_exp(cg);
+	if (distinctOps.has(AluOp.Exp) || distinctOps.has(AluOp.Erf) || distinctOps.has(AluOp.Erfc)) funcs.exp = wasm_exp(cg);
 	if (distinctOps.has(AluOp.Log)) funcs.log = wasm_log(cg);
+	if (distinctOps.has(AluOp.Erf)) funcs.erf = wasm_erf(cg, funcs.exp);
+	if (distinctOps.has(AluOp.Erfc)) funcs.erfc = wasm_erfc(cg, funcs.exp);
 	if (distinctOps.has(AluOp.Threefry2x32)) funcs.threefry2x32 = wasm_threefry2x32(cg);
 	const kernelFunc = cg.function(rep(kernel.nargs + 1, cg.i32), [], () => {
 		const gidx = cg.local.declare(cg.i32);
@@ -3633,6 +3752,8 @@ function translateExp(cg, funcs, exp, ctx) {
 		else if (op === AluOp.Atan) gen(src[0]), cg.call(funcs.atan);
 		else if (op === AluOp.Exp) gen(src[0]), cg.call(funcs.exp);
 		else if (op === AluOp.Log) gen(src[0]), cg.call(funcs.log);
+		else if (op === AluOp.Erf) gen(src[0]), cg.call(funcs.erf);
+		else if (op === AluOp.Erfc) gen(src[0]), cg.call(funcs.erfc);
 		else if (op === AluOp.Sqrt) gen(src[0]), cg.f32.sqrt();
 		else if (op === AluOp.Reciprocal) cg.f32.const(1), gen(src[0]), cg.f32.div();
 		else if (op === AluOp.Cast) {
@@ -3760,7 +3881,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-ow0Pn_6q.js");
+		const { WebGPUBackend } = await import("./webgpu-LGi2A3mS.js");
 		const importantLimits = [
 			"maxBufferSize",
 			"maxComputeInvocationsPerWorkgroup",
@@ -3813,4 +3934,5 @@ var UnsupportedOpError = class extends Error {
 };
 
 //#endregion
-export { AluExp, AluGroup, AluOp, AluVar, DEBUG, DType, Executable, FpHash, Kernel, PPrint, Reduction, ShapeTracker, SlotError, UnsupportedOpError, accessorAluExp, accessorGlobal, byteWidth, checkAxis, deepEqual, defaultDevice, devices, dtypedArray, dtypedJsArray, findPow2, getBackend, init, invertPermutation, isFloatDtype, isNumberPair, isPermutation, normalizeAxis, partitionList, prod, promoteTypes, range, recursiveFlatten, rep, runWithCache, setDebug, strip1, toposort, tuneWebgpu, union, unravelAlu, unzip2, zip, zipn };
+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-DwIAd0AG.js.map