@jax-js/jax 0.0.5 → 0.1.1

package/dist/{backend-CdcTZEOF.js → backend-CoVtc9dx.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) {
@@ -300,6 +306,7 @@ let DType = /* @__PURE__ */ function(DType$1) {
 	DType$1["Uint32"] = "uint32";
 	DType$1["Bool"] = "bool";
 	DType$1["Float16"] = "float16";
+	DType$1["Float64"] = "float64";
 	return DType$1;
 }({});
 const byteWidth = (dtype) => {
@@ -309,10 +316,11 @@ const byteWidth = (dtype) => {
 		case DType.Uint32:
 		case DType.Bool: return 4;
 		case DType.Float16: return 2;
+		case DType.Float64: return 8;
 		default: throw new TypeError(`Unknown dtype: ${dtype}`);
 	}
 };
-const isFloatDtype = (dtype) => dtype === DType.Float32 || dtype === DType.Float16;
+const isFloatDtype = (dtype) => dtype === DType.Float32 || dtype === DType.Float16 || dtype === DType.Float64;
 /**
 * Promote two dtypes to their join according to the type lattice.
 *
@@ -322,7 +330,7 @@ const isFloatDtype = (dtype) => dtype === DType.Float32 || dtype === DType.Float
 *
 * **Type lattice:**
 * ```text
-* bool -> uint32 -> int32 -> float16 -> float32
+* bool -> uint32 -> int32 -> float16 -> float32 -> float64
 *  weakType --^
 * ```
 *
@@ -344,7 +352,8 @@ function promoteTypes(dtype1, dtype2) {
 		[DType.Uint32]: 1,
 		[DType.Int32]: 2,
 		[DType.Float16]: 3,
-		[DType.Float32]: 4
+		[DType.Float32]: 4,
+		[DType.Float64]: 5
 	};
 	return rank[dtype1] > rank[dtype2] ? dtype1 : dtype2;
 }
@@ -357,6 +366,7 @@ function dtypedArray(dtype, data) {
 		case DType.Bool: return new Int32Array(buffer, byteOffset, length);
 		case DType.Uint32: return new Uint32Array(buffer, byteOffset, length);
 		case DType.Float16: return new Float16Array(buffer, byteOffset, length);
+		case DType.Float64: return new Float64Array(buffer, byteOffset, length);
 		default: throw new Error(`Unimplemented dtype: ${dtype}`);
 	}
 }
@@ -367,6 +377,7 @@ function dtypedJsArray(dtype, data) {
 		case DType.Bool: return new Int32Array(data);
 		case DType.Uint32: return new Uint32Array(data);
 		case DType.Float16: return new Float16Array(data);
+		case DType.Float64: return new Float64Array(data);
 		default: throw new Error(`Unimplemented dtype: ${dtype}`);
 	}
 }
@@ -429,6 +440,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]);
 	}
@@ -498,6 +515,9 @@ var AluExp = class AluExp {
 	static f16(value) {
 		return AluExp.const(DType.Float16, value);
 	}
+	static f64(value) {
+		return AluExp.const(DType.Float64, value);
+	}
 	not() {
 		if (this.dtype !== DType.Bool) throw new Error("not() can only be called on boolean expressions");
 		return AluExp.cmpne(this, AluExp.const(DType.Bool, true));
@@ -508,7 +528,8 @@ var AluExp = class AluExp {
 		const hasher = new FpHash();
 		hasher.update(this.op);
 		hasher.update(this.dtype);
-		hasher.update(JSON.stringify(this.arg));
+		if (this.op === AluOp.Const) hasher.update(this.arg);
+		else hasher.update(JSON.stringify(this.arg));
 		hasher.update(this.src.length);
 		for (const s of this.src) hasher.update(s);
 		this.#hash = hasher.value;
@@ -599,6 +620,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;
@@ -739,6 +766,7 @@ var AluExp = class AluExp {
 			if (op === AluOp.Mul && x === 1) return src[1 - i];
 			if (op === AluOp.Mul && x === 0) return AluExp.const(this.dtype, 0);
 			if (op === AluOp.Idiv && i === 1 && x === 1) return src[1 - i];
+			if (op === AluOp.Cmpne && src[i].dtype === DType.Bool && x === 0) return src[1 - i];
 		}
 		if ((op === AluOp.Add || op === AluOp.Sub) && src[1].op === AluOp.Mul) {
 			const [a, b] = src[1].src;
@@ -920,6 +948,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: {
@@ -938,11 +968,13 @@ var AluExp = class AluExp {
 					else if (fromType === DType.Int32) view.setInt32(0, x, true);
 					else if (fromType === DType.Uint32) view.setUint32(0, x, true);
 					else if (fromType === DType.Float16) view.setFloat16(0, x, true);
+					else if (fromType === DType.Float64) view.setFloat64(0, x, true);
 					else throw new Error(`Unsupported bitcast from ${fromType}`);
 					if (this.dtype === DType.Float32) return view.getFloat32(0, true);
 					else if (this.dtype === DType.Int32) return view.getInt32(0, true);
 					else if (this.dtype === DType.Uint32) return view.getUint32(0, true);
 					else if (this.dtype === DType.Float16) return view.getFloat16(0, true);
+					else if (this.dtype === DType.Float64) return view.getFloat64(0, true);
 					else throw new Error(`Unsupported bitcast to ${this.dtype}`);
 				}
 				default: throw new Error(`Missing implemementation for ${this.op}`);
@@ -1068,11 +1100,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;
 	}
@@ -1101,6 +1137,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";
@@ -1133,6 +1171,8 @@ const AluGroup = {
 		AluOp.Atan,
 		AluOp.Exp,
 		AluOp.Log,
+		AluOp.Erf,
+		AluOp.Erfc,
 		AluOp.Sqrt,
 		AluOp.Reciprocal,
 		AluOp.Cast,
@@ -1158,6 +1198,8 @@ const AluGroup = {
 		AluOp.Atan,
 		AluOp.Exp,
 		AluOp.Log,
+		AluOp.Erf,
+		AluOp.Erfc,
 		AluOp.Sqrt,
 		AluOp.Reciprocal
 	])
@@ -1333,6 +1375,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
@@ -2018,7 +2098,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;
@@ -2237,6 +2317,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.
 *
@@ -2277,27 +2370,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);
@@ -2325,11 +2406,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();
@@ -2366,41 +2442,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();
 	});
 }
@@ -2410,7 +2463,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);
@@ -2446,35 +2499,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,
@@ -2556,24 +2596,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();
@@ -2627,6 +2659,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,
@@ -2837,6 +2937,7 @@ var CodeGenerator = class {
 	local;
 	i32;
 	f32;
+	f64;
 	v128;
 	i32x4;
 	f32x4;
@@ -2856,6 +2957,7 @@ var CodeGenerator = class {
 		this.local = new Local(this);
 		this.i32 = new I32(this);
 		this.f32 = new F32(this);
+		this.f64 = new F64(this);
 		this.v128 = new V128(this);
 		this.i32x4 = new I32x4(this);
 		this.f32x4 = new F32x4(this);
@@ -3242,6 +3344,8 @@ var I32 = class {
 	ne = BINARY_OP("ne", 71, "i32", "i32", "i32");
 	trunc_f32_s = UNARY_OP("trunc_f32_s", 168, "f32", "i32");
 	trunc_f32_u = UNARY_OP("trunc_f32_u", 169, "f32", "i32");
+	trunc_f64_s = UNARY_OP("trunc_f64_s", 170, "f64", "i32");
+	trunc_f64_u = UNARY_OP("trunc_f64_u", 171, "f64", "i32");
 	load = LOAD_OP("load", 40, "i32");
 	load8_s = LOAD_OP("load8_s", 44, "i32");
 	load8_u = LOAD_OP("load8_u", 45, "i32");
@@ -3253,6 +3357,8 @@ var I32 = class {
 	reinterpret_f32 = UNARY_OP("reinterpret_f32", 188, "f32", "i32");
 	trunc_sat_f32_s = UNARY_OP("trunc_sat_f32_s", [252, 0], "f32", "i32");
 	trunc_sat_f32_u = UNARY_OP("trunc_sat_f32_u", [252, 1], "f32", "i32");
+	trunc_sat_f64_s = UNARY_OP("trunc_sat_f64_s", [252, 2], "f64", "i32");
+	trunc_sat_f64_u = UNARY_OP("trunc_sat_f64_u", [252, 3], "f64", "i32");
 };
 var F32 = class {
 	constructor(cg) {
@@ -3272,6 +3378,8 @@ var F32 = class {
 		for (let i = 0; i < 4; i++) this.cg._emit(bytes[i]);
 		this.cg._push(this);
 	}
+	load = LOAD_OP("load", 42, "f32");
+	store = STORE_OP("store", 56, "f32");
 	eq = BINARY_OP("eq", 91, "f32", "f32", "i32");
 	ne = BINARY_OP("ne", 92, "f32", "f32", "i32");
 	lt = BINARY_OP("lt", 93, "f32", "f32", "i32");
@@ -3294,10 +3402,53 @@ var F32 = class {
 	copysign = BINARY_OP("copysign", 152, "f32", "f32", "f32");
 	convert_i32_s = UNARY_OP("convert_i32_s", 178, "i32", "f32");
 	convert_i32_u = UNARY_OP("convert_i32_u", 179, "i32", "f32");
-	load = LOAD_OP("load", 42, "f32");
-	store = STORE_OP("store", 56, "f32");
+	demote_f64 = UNARY_OP("demote_f64", 182, "f64", "f32");
 	reinterpret_i32 = UNARY_OP("reinterpret_i32", 190, "i32", "f32");
 };
+var F64 = class {
+	constructor(cg) {
+		this.cg = cg;
+	}
+	get typeId() {
+		return 124;
+	}
+	get name() {
+		return "f64";
+	}
+	const(f) {
+		this.cg._emit(68);
+		const buffer = /* @__PURE__ */ new ArrayBuffer(8);
+		new DataView(buffer).setFloat64(0, f, true);
+		const bytes = new Uint8Array(buffer);
+		for (let i = 0; i < 8; i++) this.cg._emit(bytes[i]);
+		this.cg._push(this);
+	}
+	load = LOAD_OP("load", 43, "f64");
+	store = STORE_OP("store", 57, "f64");
+	eq = BINARY_OP("eq", 97, "f64", "f64", "i32");
+	ne = BINARY_OP("ne", 98, "f64", "f64", "i32");
+	lt = BINARY_OP("lt", 99, "f64", "f64", "i32");
+	gt = BINARY_OP("gt", 100, "f64", "f64", "i32");
+	le = BINARY_OP("le", 101, "f64", "f64", "i32");
+	ge = BINARY_OP("ge", 102, "f64", "f64", "i32");
+	abs = UNARY_OP("abs", 153, "f64", "f64");
+	neg = UNARY_OP("neg", 154, "f64", "f64");
+	ceil = UNARY_OP("ceil", 155, "f64", "f64");
+	floor = UNARY_OP("floor", 156, "f64", "f64");
+	trunc = UNARY_OP("trunc", 157, "f64", "f64");
+	nearest = UNARY_OP("nearest", 158, "f64", "f64");
+	sqrt = UNARY_OP("sqrt", 159, "f64", "f64");
+	add = BINARY_OP("add", 160, "f64", "f64", "f64");
+	sub = BINARY_OP("sub", 161, "f64", "f64", "f64");
+	mul = BINARY_OP("mul", 162, "f64", "f64", "f64");
+	div = BINARY_OP("div", 163, "f64", "f64", "f64");
+	min = BINARY_OP("min", 164, "f64", "f64", "f64");
+	max = BINARY_OP("max", 165, "f64", "f64", "f64");
+	copysign = BINARY_OP("copysign", 166, "f64", "f64", "f64");
+	convert_i32_s = UNARY_OP("convert_i32_s", 183, "i32", "f64");
+	convert_i32_u = UNARY_OP("convert_i32_u", 184, "i32", "f64");
+	promote_f32 = UNARY_OP("promote_f32", 187, "f32", "f64");
+};
 function VECTOR_OP(op, vopcode, inTypes, outType) {
 	return function() {
 		for (const inType of inTypes.toReversed()) {
@@ -3501,14 +3652,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);
@@ -3548,10 +3701,10 @@ function codegenWasm(kernel) {
 				cg.local.get(acc);
 				if (re.dtype === DType.Bool) cg.i32.and();
 				else dty(cg, re.op, re.dtype).mul();
-			} else if (re.op === AluOp.Min || re.op === AluOp.Max) if (re.dtype === DType.Float32) {
+			} else if (re.op === AluOp.Min || re.op === AluOp.Max) if (isFloatDtype(re.dtype)) {
 				cg.local.get(acc);
-				if (re.op === AluOp.Min) cg.f32.min();
-				else cg.f32.max();
+				if (re.op === AluOp.Min) dtyF(cg, re.op, re.dtype).min();
+				else dtyF(cg, re.op, re.dtype).max();
 			} else if ([
 				DType.Int32,
 				DType.Uint32,
@@ -3613,27 +3766,30 @@ function translateExp(cg, funcs, exp, ctx) {
 			else if (op === AluOp.Sub) dty(cg, op, dtype).sub();
 			else if (op === AluOp.Mul) if (dtype === DType.Bool) cg.i32.and();
 			else dty(cg, op, dtype).mul();
-			else if (op === AluOp.Idiv) if (dtype === DType.Float32) cg.f32.div(), cg.f32.trunc();
-			else if (dtype === DType.Uint32) cg.i32.div_u();
+			else if (op === AluOp.Idiv) if (isFloatDtype(dtype)) {
+				dtyF(cg, op, dtype).div();
+				dtyF(cg, op, dtype).trunc();
+			} else if (dtype === DType.Uint32) cg.i32.div_u();
 			else if (dtype === DType.Int32) cg.i32.div_s();
 			else throw new UnsupportedOpError(op, dtype, "wasm");
-			else if (op === AluOp.Mod) if (dtype === DType.Float32) {
-				const a = cg.local.declare(cg.f32);
-				const b = cg.local.declare(cg.f32);
+			else if (op === AluOp.Mod) if (isFloatDtype(dtype)) {
+				const dt = dtyF(cg, op, dtype);
+				const a = cg.local.declare(dt);
+				const b = cg.local.declare(dt);
 				cg.local.set(b);
 				cg.local.tee(a);
 				cg.local.get(a);
 				cg.local.get(b);
-				cg.f32.div();
-				cg.f32.trunc();
+				dt.div();
+				dt.trunc();
 				cg.local.get(b);
-				cg.f32.mul();
-				cg.f32.sub();
+				dt.mul();
+				dt.sub();
 			} else if (dtype === DType.Uint32) cg.i32.rem_u();
 			else if (dtype === DType.Int32) cg.i32.rem_s();
 			else throw new UnsupportedOpError(op, dtype, "wasm");
-			else if (op === AluOp.Min || op === AluOp.Max) if (dtype === DType.Float32) if (op === AluOp.Min) cg.f32.min();
-			else cg.f32.max();
+			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) {
 				const a = cg.local.declare(cg.i32);
 				const b = cg.local.declare(cg.i32);
@@ -3650,52 +3806,74 @@ function translateExp(cg, funcs, exp, ctx) {
 			} else throw new UnsupportedOpError(op, dtype, "wasm");
 			else if (op === AluOp.Cmplt) {
 				const srcDtype = src[0].dtype;
-				if (srcDtype === DType.Float32) cg.f32.lt();
+				if (isFloatDtype(srcDtype)) dtyF(cg, op, srcDtype).lt();
 				else if (srcDtype === DType.Int32) cg.i32.lt_s();
 				else if (srcDtype === DType.Uint32) cg.i32.lt_u();
 				else throw new UnsupportedOpError(op, dtype, "wasm");
 			} else if (op === AluOp.Cmpne) dty(cg, op, src[0].dtype).ne();
 			else throw new UnsupportedOpError(op, dtype, "wasm");
-		} else if (AluGroup.Unary.has(op)) if (op === AluOp.Sin) gen(src[0]), cg.call(funcs.sin);
-		else if (op === AluOp.Cos) gen(src[0]), cg.call(funcs.cos);
-		else if (op === AluOp.Asin) gen(src[0]), cg.call(funcs.asin);
-		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.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) {
-			gen(src[0]);
-			const dtype0 = src[0].dtype;
-			const i32repr = dtype0 === DType.Int32 || dtype0 === DType.Uint32 || dtype0 === DType.Bool;
-			if (dtype === DType.Int32) if (dtype0 === DType.Float32) cg.i32.trunc_sat_f32_s();
-			else if (i32repr);
-			else throw new UnsupportedOpError(op, dtype, "wasm", dtype0);
-			else if (dtype === DType.Uint32) if (dtype0 === DType.Float32) cg.i32.trunc_sat_f32_u();
-			else if (i32repr);
-			else throw new UnsupportedOpError(op, dtype, "wasm", dtype0);
-			else if (dtype === DType.Float32) if (dtype0 === DType.Float32);
-			else if (dtype0 === DType.Int32 || dtype0 === DType.Bool) cg.f32.convert_i32_s();
-			else if (dtype0 === DType.Uint32) cg.f32.convert_i32_u();
-			else throw new UnsupportedOpError(op, dtype, "wasm", dtype0);
-			else if (dtype === DType.Bool) if (dtype0 === DType.Bool);
-			else if (i32repr) cg.i32.const(0), cg.i32.ne();
-			else if (dtype0 === DType.Float32) cg.f32.const(0), cg.f32.ne();
-			else throw new UnsupportedOpError(op, dtype, "wasm", dtype0);
-			else throw new UnsupportedOpError(op, dtype, "wasm");
-		} else if (op === AluOp.Bitcast) {
-			gen(src[0]);
-			const dtype0 = src[0].dtype;
-			const i32repr = dtype0 === DType.Int32 || dtype0 === DType.Uint32;
-			if (dtype === DType.Int32 || dtype === DType.Uint32) if (dtype0 === DType.Float32) cg.i32.reinterpret_f32();
-			else if (i32repr);
-			else throw new UnsupportedOpError(op, dtype, "wasm", dtype0);
-			else if (dtype === DType.Float32) if (i32repr) cg.f32.reinterpret_i32();
-			else if (dtype0 === DType.Float32);
-			else throw new UnsupportedOpError(op, dtype, "wasm", dtype0);
-			else throw new UnsupportedOpError(op, dtype, "wasm");
-		} else throw new UnsupportedOpError(op, dtype, "wasm");
-		else if (op === AluOp.Where) {
+		} else if (AluGroup.Unary.has(op)) {
+			const callFuncF32 = (func) => {
+				if (dtype !== DType.Float32) if (dtype === DType.Float64) cg.f32.demote_f64();
+				else throw new UnsupportedOpError(op, dtype, "wasm");
+				cg.call(func);
+				if (dtype === DType.Float64) cg.f64.promote_f32();
+			};
+			if (op === AluOp.Sin) gen(src[0]), callFuncF32(funcs.sin);
+			else if (op === AluOp.Cos) gen(src[0]), callFuncF32(funcs.cos);
+			else if (op === AluOp.Asin) gen(src[0]), callFuncF32(funcs.asin);
+			else if (op === AluOp.Atan) gen(src[0]), callFuncF32(funcs.atan);
+			else if (op === AluOp.Exp) gen(src[0]), callFuncF32(funcs.exp);
+			else if (op === AluOp.Log) gen(src[0]), callFuncF32(funcs.log);
+			else if (op === AluOp.Erf) gen(src[0]), callFuncF32(funcs.erf);
+			else if (op === AluOp.Erfc) gen(src[0]), callFuncF32(funcs.erfc);
+			else if (op === AluOp.Sqrt) gen(src[0]), dtyF(cg, op, dtype).sqrt();
+			else if (op === AluOp.Reciprocal) {
+				const dt = dtyF(cg, op, dtype);
+				dt.const(1), gen(src[0]), dt.div();
+			} else if (op === AluOp.Cast) {
+				gen(src[0]);
+				const dtype0 = src[0].dtype;
+				const i32repr = dtype0 === DType.Int32 || dtype0 === DType.Uint32 || dtype0 === DType.Bool;
+				if (dtype === DType.Int32) if (dtype0 === DType.Float32) cg.i32.trunc_sat_f32_s();
+				else if (dtype0 === DType.Float64) cg.i32.trunc_sat_f64_s();
+				else if (i32repr);
+				else throw new UnsupportedOpError(op, dtype, "wasm", dtype0);
+				else if (dtype === DType.Uint32) if (dtype0 === DType.Float32) cg.i32.trunc_sat_f32_u();
+				else if (dtype0 === DType.Float64) cg.i32.trunc_sat_f64_u();
+				else if (i32repr);
+				else throw new UnsupportedOpError(op, dtype, "wasm", dtype0);
+				else if (dtype === DType.Float32) if (dtype0 === DType.Float32);
+				else if (dtype0 === DType.Float64) cg.f32.demote_f64();
+				else if (dtype0 === DType.Int32 || dtype0 === DType.Bool) cg.f32.convert_i32_s();
+				else if (dtype0 === DType.Uint32) cg.f32.convert_i32_u();
+				else throw new UnsupportedOpError(op, dtype, "wasm", dtype0);
+				else if (dtype === DType.Float64) if (dtype0 === DType.Float32) cg.f64.promote_f32();
+				else if (dtype0 === DType.Float64);
+				else if (dtype0 === DType.Int32 || dtype0 === DType.Bool) cg.f64.convert_i32_s();
+				else if (dtype0 === DType.Uint32) cg.f64.convert_i32_u();
+				else throw new UnsupportedOpError(op, dtype, "wasm", dtype0);
+				else if (dtype === DType.Bool) if (dtype0 === DType.Bool);
+				else if (i32repr) cg.i32.const(0), cg.i32.ne();
+				else if (dtype0 === DType.Float32) cg.f32.const(0), cg.f32.ne();
+				else if (dtype0 === DType.Float64) cg.f64.const(0), cg.f64.ne();
+				else throw new UnsupportedOpError(op, dtype, "wasm", dtype0);
+				else throw new UnsupportedOpError(op, dtype, "wasm");
+			} else if (op === AluOp.Bitcast) {
+				gen(src[0]);
+				const dtype0 = src[0].dtype;
+				if (dtype !== dtype0) {
+					const i32repr = dtype0 === DType.Int32 || dtype0 === DType.Uint32;
+					if (dtype === DType.Int32 || dtype === DType.Uint32) if (dtype0 === DType.Float32) cg.i32.reinterpret_f32();
+					else if (i32repr);
+					else throw new UnsupportedOpError(op, dtype, "wasm", dtype0);
+					else if (dtype === DType.Float32) if (i32repr) cg.f32.reinterpret_i32();
+					else if (dtype0 === DType.Float32);
+					else throw new UnsupportedOpError(op, dtype, "wasm", dtype0);
+					else throw new UnsupportedOpError(op, dtype, "wasm");
+				}
+			} else throw new UnsupportedOpError(op, dtype, "wasm");
+		} else if (op === AluOp.Where) {
 			gen(src[1]);
 			gen(src[2]);
 			gen(src[0]);
@@ -3740,12 +3918,20 @@ function translateExp(cg, funcs, exp, ctx) {
 function dty(cg, op, dtype) {
 	switch (dtype) {
 		case DType.Float32: return cg.f32;
+		case DType.Float64: return cg.f64;
 		case DType.Int32:
 		case DType.Uint32:
 		case DType.Bool: return cg.i32;
 		default: throw new UnsupportedOpError(op, dtype, "wasm");
 	}
 }
+function dtyF(cg, op, dtype) {
+	switch (dtype) {
+		case DType.Float32: return cg.f32;
+		case DType.Float64: return cg.f64;
+		default: throw new UnsupportedOpError(op, dtype, "wasm");
+	}
+}
 
 //#endregion
 //#region src/backend.ts
@@ -3754,10 +3940,10 @@ const devices = [
 	"wasm",
 	"webgpu"
 ];
-let defaultBackend = "wasm";
 const initializedBackends = /* @__PURE__ */ new Map();
 initializedBackends.set("cpu", new CpuBackend());
-initializedBackends.set("wasm", new WasmBackend());
+if (typeof WebAssembly !== "undefined") initializedBackends.set("wasm", new WasmBackend());
+let defaultBackend = initializedBackends.has("wasm") ? "wasm" : "cpu";
 /** Configure the default device for arrays. */
 function defaultDevice(device) {
 	if (device !== void 0) if (initializedBackends.has(device)) defaultBackend = device;
@@ -3784,12 +3970,14 @@ async function init(...devicesToInit) {
 /** Create a backend, if available. Internal function called by `init()`. */
 async function createBackend(device) {
 	if (device === "cpu") return new CpuBackend();
-	else if (device === "wasm") return new WasmBackend();
-	else if (device === "webgpu") {
+	else if (device === "wasm") {
+		if (typeof WebAssembly === "undefined") return null;
+		return new WasmBackend();
+	} else if (device === "webgpu") {
 		if (!navigator.gpu) return null;
 		const adapter = await navigator.gpu.requestAdapter({ powerPreference: "high-performance" });
 		if (!adapter) return null;
-		const { WebGPUBackend } = await import("./webgpu-CM-xNYzW.js");
+		const { WebGPUBackend } = await import("./webgpu-B3UVme6n.js");
 		const importantLimits = [
 			"maxBufferSize",
 			"maxComputeInvocationsPerWorkgroup",
@@ -3842,4 +4030,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, generalBroadcast, 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-CoVtc9dx.js.map