@jax-js/jax 0.0.3 → 0.0.5

package/dist/{backend-D2C4MJRP.cjs → backend-yEU0L_ig.cjs}

@@ -36,7 +36,22 @@ var PPrint = class PPrint {
 //#endregion
 //#region src/utils.ts
 /** @file Generic programming utilities with no dependencies on library code. */
-const DEBUG = 3;
+let DEBUG = 0;
+/**
+* Set the debug level for verbose logging.
+*
+* 1. JIT compile logs
+* 2. Shader code
+* 3. Expressions and metadata
+* 4. JIT programs, tuning details
+* 5. Most verbose operation traces
+*
+* This is an experimental API and may change in behavior. Do not rely on this
+* in production.
+*/
+function setDebug(level) {
+	DEBUG = level;
+}
 function unzip2(pairs) {
 	const lst1 = [];
 	const lst2 = [];
@@ -110,9 +125,23 @@ function isNumberPair(x) {
 }
 /** Check an axis against number of dimensions, and resolve negative axes. */
 function checkAxis(axis, ndim) {
-	if (axis < -ndim || axis >= ndim) throw new Error(`Invalid axis ${axis} for array of ${ndim} dimensions`);
+	if (axis < -ndim || axis >= ndim) throw new Error(`Axis ${axis} out of bounds for array of dimension ${ndim}`);
 	return axis < 0 ? axis + ndim : axis;
 }
+/** Normalize common axis argument for functions, defaulting to all axes. */
+function normalizeAxis(axis, ndim) {
+	if (axis === null) return range(ndim);
+	else if (typeof axis === "number") return [checkAxis(axis, ndim)];
+	else {
+		const seen = /* @__PURE__ */ new Set();
+		for (const a of axis) {
+			const ca = checkAxis(a, ndim);
+			if (seen.has(ca)) throw new Error(`Duplicate axis ${ca} passed to function`);
+			seen.add(ca);
+		}
+		return [...seen].sort();
+	}
+}
 function range(start, stop, step = 1) {
 	if (stop === void 0) {
 		stop = start;
@@ -178,6 +207,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);
@@ -187,6 +245,7 @@ function strip1(str) {
 	if (str[0] === "(" && str[str.length - 1] === ")") return str.slice(1, -1);
 	return str;
 }
+const _stagingbuf = /* @__PURE__ */ new DataView(/* @__PURE__ */ new ArrayBuffer(8));
 /**
 * Polynomial hashes modulo p are good at avoiding collisions in expectation.
 * Probability-wise, it's good enough to be used for something like
@@ -201,22 +260,26 @@ var FpHash = class FpHash {
 		const modulus = 3189051996290219n;
 		this.value = (this.value * base + x) % modulus;
 	}
-	update(...values) {
-		for (const x of values) if (typeof x === "string") for (const c of x) this.#update(BigInt(199 + c.charCodeAt(0)));
-		else if (typeof x === "number") if (Number.isInteger(x)) this.#update(68265653n ^ BigInt(x));
+	update(x) {
+		if (typeof x === "string") {
+			this.#update(BigInt(x.length));
+			for (let i = 0; i < x.length; i++) this.#update(BigInt(199 + x.charCodeAt(i)));
+		} else if (typeof x === "number") if (Number.isInteger(x)) this.#update(68265653n ^ BigInt(x));
 		else {
-			const ar = new Float64Array([x]);
-			this.#update(new DataView(ar.buffer).getBigUint64(0, true));
+			_stagingbuf.setFloat64(0, x, true);
+			this.#update(_stagingbuf.getBigUint64(0, true));
 		}
 		else if (typeof x === "boolean") this.#update(x ? 69069841n : 63640693n);
 		else if (typeof x === "bigint") this.#update(x ^ 71657401n);
 		else if (x === null) this.#update(37832657n);
 		else if (x === void 0) this.#update(18145117n);
-		else if (typeof x === "object" && "hash" in x) x.hash(this);
+		else x.hash(this);
 		return this;
 	}
 	static hash(...values) {
-		return new FpHash().update(...values).value;
+		const h = new FpHash();
+		for (const x of values) h.update(x);
+		return h.value;
 	}
 };
 /** Run a function while caching it inline inside a `Map`. */
@@ -251,6 +314,41 @@ const byteWidth = (dtype) => {
 	}
 };
 const isFloatDtype = (dtype) => dtype === DType.Float32 || dtype === DType.Float16;
+/**
+* Promote two dtypes to their join according to the type lattice.
+*
+* When performing operations between arrays of different types, we need to
+* promote both operands to a common type that can represent values from both
+* input types. This follows JAX's type promotion rules.
+*
+* **Type lattice:**
+* ```text
+* bool -> uint32 -> int32 -> float16 -> float32
+*  weakType --^
+* ```
+*
+* `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`
+* - `promoteTypes(uint32, int32) → int32`
+* - `promoteTypes(int32, float16) → float16`
+* - `promoteTypes(float16, float32) → float32`
+* - `promoteTypes(uint32, float32) → float32`
+*/
+function promoteTypes(dtype1, dtype2) {
+	if (dtype1 === dtype2) return dtype1;
+	const rank = {
+		[DType.Bool]: 0,
+		[DType.Uint32]: 1,
+		[DType.Int32]: 2,
+		[DType.Float16]: 3,
+		[DType.Float32]: 4
+	};
+	return rank[dtype1] > rank[dtype2] ? dtype1 : dtype2;
+}
 function dtypedArray(dtype, data) {
 	const { buffer, byteLength, byteOffset } = data;
 	const length = byteLength / byteWidth(dtype);
@@ -320,6 +418,12 @@ var AluExp = class AluExp {
 	static cos(a) {
 		return new AluExp(AluOp.Cos, a.dtype, [a]);
 	}
+	static asin(a) {
+		return new AluExp(AluOp.Asin, a.dtype, [a]);
+	}
+	static atan(a) {
+		return new AluExp(AluOp.Atan, a.dtype, [a]);
+	}
 	static exp(a) {
 		return new AluExp(AluOp.Exp, a.dtype, [a]);
 	}
@@ -403,8 +507,11 @@ var AluExp = class AluExp {
 	getHash() {
 		if (this.#hash !== void 0) return this.#hash;
 		const hasher = new FpHash();
-		hasher.update(this.op, this.dtype, JSON.stringify(this.arg));
-		hasher.update(this.src.length, ...this.src);
+		hasher.update(this.op);
+		hasher.update(this.dtype);
+		hasher.update(JSON.stringify(this.arg));
+		hasher.update(this.src.length);
+		for (const s of this.src) hasher.update(s);
 		this.#hash = hasher.value;
 		return this.#hash;
 	}
@@ -476,10 +583,16 @@ var AluExp = class AluExp {
 				ret = [Math.max(src[0].min, src[1].min), Math.max(src[0].max, src[1].max)];
 				break;
 			case AluOp.Sin:
-				ret = [Math.sin(src[0].min), Math.sin(src[0].max)];
+				ret = [-1, 1];
 				break;
 			case AluOp.Cos:
-				ret = [Math.cos(src[0].min), Math.cos(src[0].max)];
+				ret = [-1, 1];
+				break;
+			case AluOp.Asin:
+				ret = [-Math.PI / 2, Math.PI / 2];
+				break;
+			case AluOp.Atan:
+				ret = [-Math.PI / 2, Math.PI / 2];
 				break;
 			case AluOp.Exp:
 				ret = [Math.exp(src[0].min), Math.exp(src[0].max)];
@@ -595,11 +708,12 @@ var AluExp = class AluExp {
 	simplify(cache = /* @__PURE__ */ new Map()) {
 		if (this.#simplified !== void 0) return this.#simplified;
 		const hash = this.getHash();
-		if (cache.has(hash)) return this.#simplified = cache.get(hash);
+		const prevCachedValue = cache.get(hash);
+		if (prevCachedValue !== void 0) return this.#simplified = prevCachedValue;
 		const simplified = this.#simplifyInner(cache);
 		const simplifiedHash = simplified.getHash();
-		if (cache.has(simplifiedHash)) {
-			const prevSimplified = cache.get(simplifiedHash);
+		const prevSimplified = cache.get(simplifiedHash);
+		if (prevSimplified !== void 0) {
 			cache.set(hash, prevSimplified);
 			this.#simplified = prevSimplified;
 			return prevSimplified;
@@ -803,6 +917,8 @@ var AluExp = class AluExp {
 			switch (this.op) {
 				case AluOp.Sin: return Math.sin(x);
 				case AluOp.Cos: return Math.cos(x);
+				case AluOp.Asin: return Math.asin(x);
+				case AluOp.Atan: return Math.atan(x);
 				case AluOp.Exp: return Math.exp(x);
 				case AluOp.Log: return Math.log(x);
 				case AluOp.Sqrt: return Math.sqrt(x);
@@ -982,6 +1098,8 @@ let AluOp = /* @__PURE__ */ function(AluOp$1) {
 	AluOp$1["Max"] = "Max";
 	AluOp$1["Sin"] = "Sin";
 	AluOp$1["Cos"] = "Cos";
+	AluOp$1["Asin"] = "Asin";
+	AluOp$1["Atan"] = "Atan";
 	AluOp$1["Exp"] = "Exp";
 	AluOp$1["Log"] = "Log";
 	AluOp$1["Sqrt"] = "Sqrt";
@@ -1012,6 +1130,8 @@ const AluGroup = {
 	Unary: new Set([
 		AluOp.Sin,
 		AluOp.Cos,
+		AluOp.Asin,
+		AluOp.Atan,
 		AluOp.Exp,
 		AluOp.Log,
 		AluOp.Sqrt,
@@ -1035,6 +1155,8 @@ const AluGroup = {
 	RequiredFloat: new Set([
 		AluOp.Sin,
 		AluOp.Cos,
+		AluOp.Asin,
+		AluOp.Atan,
 		AluOp.Exp,
 		AluOp.Log,
 		AluOp.Sqrt,
@@ -1066,7 +1188,7 @@ var Kernel = class {
 		this.exp = exp.simplify();
 	}
 	hash(state) {
-		state.update(this.nargs, this.size, this.exp, this.reduction);
+		state.update(this.nargs).update(this.size).update(this.exp).update(this.reduction);
 	}
 	pprint() {
 		let details = PPrint.pp(`exp = ${this.exp}`);
@@ -1112,7 +1234,7 @@ var Reduction = class {
 		this.epilogue = epilogue.simplify();
 	}
 	hash(state) {
-		state.update(this.dtype, this.op, this.size, this.epilogue);
+		state.update(this.dtype).update(this.op).update(this.size).update(this.epilogue);
 	}
 	toString() {
 		return `${this.op}{${this.size}} -> ${this.epilogue}`;
@@ -2284,78 +2406,92 @@ function wasm_log(cg) {
 	});
 }
 /**
-* Approximate sin(x).
+* Common helper to approximate sin(x) and cos(x).
 *
 * Method: reduce to y in [-π, π], then quadrant via q = round(y/(π/2))
-*         z = y - q*(π/2); use odd polynomial on z:
+*         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)
+*/
+function _sincos(cg) {
+	const y = cg.local.declare(cg.f32);
+	const qf = cg.local.declare(cg.f32);
+	const q = cg.local.declare(cg.i32);
+	const z = cg.local.declare(cg.f32);
+	const z2 = cg.local.declare(cg.f32);
+	const sz = cg.local.declare(cg.f32);
+	const cz = cg.local.declare(cg.f32);
+	cg.local.get(0);
+	cg.local.get(0);
+	cg.f32.const(1 / (2 * Math.PI));
+	cg.f32.mul();
+	cg.f32.nearest();
+	cg.local.tee(qf);
+	cg.f32.const(2 * Math.PI);
+	cg.f32.mul();
+	cg.f32.sub();
+	cg.local.set(y);
+	cg.local.get(y);
+	cg.f32.const(2 / Math.PI);
+	cg.f32.mul();
+	cg.f32.nearest();
+	cg.local.tee(qf);
+	cg.i32.trunc_f32_s();
+	cg.local.set(q);
+	cg.local.get(y);
+	cg.local.get(qf);
+	cg.f32.const(Math.PI / 2);
+	cg.f32.mul();
+	cg.f32.sub();
+	cg.local.tee(z);
+	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();
+	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();
+	cg.local.set(cz);
+	return {
+		q,
+		sz,
+		cz
+	};
+}
+/**
+* Approximate sin(x).
+*
+* Quadrant mapping: k=q mod 4: 0: +sz, 1: +cz, 2: -sz, 3: -cz
 */
 function wasm_sin(cg) {
 	return cg.function([cg.f32], [cg.f32], () => {
-		const y = cg.local.declare(cg.f32);
-		const qf = cg.local.declare(cg.f32);
-		const q = cg.local.declare(cg.i32);
-		const z = cg.local.declare(cg.f32);
-		const z2 = cg.local.declare(cg.f32);
-		const sz = cg.local.declare(cg.f32);
-		const cz = cg.local.declare(cg.f32);
+		const { q, sz, cz } = _sincos(cg);
 		const mag = cg.local.declare(cg.f32);
-		cg.local.get(0);
-		cg.local.get(0);
-		cg.f32.const(1 / (2 * Math.PI));
-		cg.f32.mul();
-		cg.f32.nearest();
-		cg.local.tee(qf);
-		cg.f32.const(2 * Math.PI);
-		cg.f32.mul();
-		cg.f32.sub();
-		cg.local.set(y);
-		cg.local.get(y);
-		cg.f32.const(2 / Math.PI);
-		cg.f32.mul();
-		cg.f32.nearest();
-		cg.local.tee(qf);
-		cg.i32.trunc_f32_s();
-		cg.local.set(q);
-		cg.local.get(y);
-		cg.local.get(qf);
-		cg.f32.const(Math.PI / 2);
-		cg.f32.mul();
-		cg.f32.sub();
-		cg.local.tee(z);
-		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();
-		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();
-		cg.local.set(cz);
 		cg.local.get(cz);
 		cg.local.get(sz);
 		cg.local.get(q);
@@ -2374,75 +2510,12 @@ function wasm_sin(cg) {
 /**
 * Approximate cos(x).
 *
-* Same reduction as sinf, then quadrant mapping:
-* k=q mod 4: 0: +cz, 1: -sz, 2: -cz, 3: +sz
+* Quadrant mapping: k=q mod 4: 0: +cz, 1: -sz, 2: -cz, 3: +sz
 */
 function wasm_cos(cg) {
 	return cg.function([cg.f32], [cg.f32], () => {
-		const y = cg.local.declare(cg.f32);
-		const qf = cg.local.declare(cg.f32);
-		const q = cg.local.declare(cg.i32);
-		const z = cg.local.declare(cg.f32);
-		const z2 = cg.local.declare(cg.f32);
-		const sz = cg.local.declare(cg.f32);
-		const cz = cg.local.declare(cg.f32);
+		const { q, sz, cz } = _sincos(cg);
 		const mag = cg.local.declare(cg.f32);
-		cg.local.get(0);
-		cg.local.get(0);
-		cg.f32.const(1 / (2 * Math.PI));
-		cg.f32.mul();
-		cg.f32.nearest();
-		cg.local.tee(qf);
-		cg.f32.const(2 * Math.PI);
-		cg.f32.mul();
-		cg.f32.sub();
-		cg.local.set(y);
-		cg.local.get(y);
-		cg.f32.const(2 / Math.PI);
-		cg.f32.mul();
-		cg.f32.nearest();
-		cg.local.tee(qf);
-		cg.i32.trunc_f32_s();
-		cg.local.set(q);
-		cg.local.get(y);
-		cg.local.get(qf);
-		cg.f32.const(Math.PI / 2);
-		cg.f32.mul();
-		cg.f32.sub();
-		cg.local.tee(z);
-		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();
-		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();
-		cg.local.set(cz);
 		cg.local.get(sz);
 		cg.local.get(cz);
 		cg.local.get(q);
@@ -2460,6 +2533,100 @@ function wasm_cos(cg) {
 		cg.select();
 	});
 }
+/** Helper function for approximating arctan(x).  */
+function _atan(cg) {
+	const x = cg.local.declare(cg.f32);
+	const abs_x = cg.local.declare(cg.f32);
+	const z = cg.local.declare(cg.f32);
+	const z2 = cg.local.declare(cg.f32);
+	const p = cg.local.declare(cg.f32);
+	cg.local.set(x);
+	cg.local.get(x);
+	cg.f32.abs();
+	cg.local.set(abs_x);
+	cg.f32.const(1);
+	cg.local.get(abs_x);
+	cg.f32.div();
+	cg.local.get(abs_x);
+	cg.local.get(abs_x);
+	cg.f32.const(1);
+	cg.f32.ge();
+	cg.select();
+	cg.local.set(z);
+	cg.local.get(z);
+	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();
+	cg.f32.div();
+	cg.local.get(z);
+	cg.f32.mul();
+	cg.local.set(p);
+	cg.f32.const(Math.PI / 2);
+	cg.local.get(p);
+	cg.f32.sub();
+	cg.local.get(p);
+	cg.local.get(abs_x);
+	cg.f32.const(1);
+	cg.f32.ge();
+	cg.select();
+	cg.local.get(x);
+	cg.f32.copysign();
+}
+/**
+* Approximate atan(x).
+*
+* Method: if |x| < 1, use rational approximation: atan(x) ≈ x * P(x^2) / Q(x^2)
+*         where P(u) = A0 + A1*u + A2*u^2 (degree 2)
+*               Q(u) = 1 + B1*u + B2*u^2 (degree 2)
+*         if |x| >= 1, use: atan(x) = sign(x)*π/2 - atan(1/x)
+*         (fitted coefficients, max error ~5e-7 on [0,1])
+*/
+function wasm_atan(cg) {
+	return cg.function([cg.f32], [cg.f32], () => {
+		cg.local.get(0);
+		_atan(cg);
+	});
+}
+/**
+* Approximate asin(x).
+*
+* Method: asin(x) = 2 * atan(x / (1 + sqrt(1 - x^2)))
+*/
+function wasm_asin(cg) {
+	return cg.function([cg.f32], [cg.f32], () => {
+		cg.local.get(0);
+		cg.f32.const(1);
+		cg.local.get(0);
+		cg.local.get(0);
+		cg.f32.mul();
+		cg.f32.sub();
+		cg.f32.sqrt();
+		cg.f32.const(1);
+		cg.f32.add();
+		cg.f32.div();
+		_atan(cg);
+		cg.f32.const(2);
+		cg.f32.mul();
+	});
+}
 /**
 * Threefry2x32 pseudorandom number generator.
 *
@@ -3339,6 +3506,8 @@ function codegenWasm(kernel) {
 	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.Log)) funcs.log = wasm_log(cg);
 	if (distinctOps.has(AluOp.Threefry2x32)) funcs.threefry2x32 = wasm_threefry2x32(cg);
@@ -3490,6 +3659,8 @@ function translateExp(cg, funcs, exp, ctx) {
 			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();
@@ -3588,10 +3759,11 @@ let defaultBackend = "wasm";
 const initializedBackends = /* @__PURE__ */ new Map();
 initializedBackends.set("cpu", new CpuBackend());
 initializedBackends.set("wasm", new WasmBackend());
-/** Set the default device backend (must be initialized). */
-function setDevice(device) {
-	if (initializedBackends.has(device)) defaultBackend = device;
+/** Configure the default device for arrays. */
+function defaultDevice(device) {
+	if (device !== void 0) if (initializedBackends.has(device)) defaultBackend = device;
 	else throw new Error(`Backend not initialized: ${device}`);
+	return defaultBackend;
 }
 /**
 * Initialize `jax-js` library backends.
@@ -3618,7 +3790,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-fqhx41TC.cjs"));
+		const { WebGPUBackend } = await Promise.resolve().then(() => require("./webgpu-CNOpiO5T.cjs"));
 		const importantLimits = [
 			"maxBufferSize",
 			"maxComputeInvocationsPerWorkgroup",
@@ -3785,6 +3957,12 @@ Object.defineProperty(exports, 'deepEqual', {
     return deepEqual;
   }
 });
+Object.defineProperty(exports, 'defaultDevice', {
+  enumerable: true,
+  get: function () {
+    return defaultDevice;
+  }
+});
 Object.defineProperty(exports, 'devices', {
   enumerable: true,
   get: function () {
@@ -3809,6 +3987,12 @@ Object.defineProperty(exports, 'findPow2', {
     return findPow2;
   }
 });
+Object.defineProperty(exports, 'generalBroadcast', {
+  enumerable: true,
+  get: function () {
+    return generalBroadcast;
+  }
+});
 Object.defineProperty(exports, 'getBackend', {
   enumerable: true,
   get: function () {
@@ -3845,6 +4029,12 @@ Object.defineProperty(exports, 'isPermutation', {
     return isPermutation;
   }
 });
+Object.defineProperty(exports, 'normalizeAxis', {
+  enumerable: true,
+  get: function () {
+    return normalizeAxis;
+  }
+});
 Object.defineProperty(exports, 'partitionList', {
   enumerable: true,
   get: function () {
@@ -3857,6 +4047,12 @@ Object.defineProperty(exports, 'prod', {
     return prod;
   }
 });
+Object.defineProperty(exports, 'promoteTypes', {
+  enumerable: true,
+  get: function () {
+    return promoteTypes;
+  }
+});
 Object.defineProperty(exports, 'range', {
   enumerable: true,
   get: function () {
@@ -3881,10 +4077,10 @@ Object.defineProperty(exports, 'runWithCache', {
     return runWithCache;
   }
 });
-Object.defineProperty(exports, 'setDevice', {
+Object.defineProperty(exports, 'setDebug', {
   enumerable: true,
   get: function () {
-    return setDevice;
+    return setDebug;
   }
 });
 Object.defineProperty(exports, 'strip1', {