@jax-js/jax 0.0.1 → 0.0.2

package/dist/backend-1eVbAoaV.js

@@ -0,0 +1,1890 @@
+//#region src/pprint.ts
+/** General class for pretty-printing expressions with indentation. */
+var PPrint = class PPrint {
+	constructor(indents, lines) {
+		this.indents = indents;
+		this.lines = lines;
+	}
+	/** Add a fixed amount of indentation to each line. */
+	indent(spaces) {
+		return new PPrint(this.indents.map((i) => i + spaces), this.lines);
+	}
+	/** Concatenate pretty-printed expressions with newlines. */
+	concat(...items) {
+		return new PPrint((this.indents ?? []).concat(...items.map((i) => i.indents)), (this.lines ?? []).concat(...items.map((i) => i.lines)));
+	}
+	/** Stack one block to the right of another one, sharing 1 common line. */
+	stack(other) {
+		if (!other.lines.length) return this;
+		if (!this.lines.length) return other;
+		const indent = this.indents[this.indents.length - 1];
+		const s = this.lines[this.lines.length - 1];
+		const indentedBlock = other.indent(indent + s.length);
+		return new PPrint(this.indents.concat(indentedBlock.indents.slice(1)), this.lines.slice(0, -1).concat(s + " ".repeat(other.indents[0]) + other.lines[0], ...indentedBlock.lines.slice(1)));
+	}
+	/** Combine this block of lines into a formatted string. */
+	toString() {
+		return this.lines.map((line, i) => " ".repeat(this.indents[i]) + line).join("\n");
+	}
+	static pp(s) {
+		const lines = s.toString().split("\n");
+		return new PPrint(Array(lines.length).fill(0), lines);
+	}
+};
+
+//#endregion
+//#region src/utils.ts
+/** @file Generic programming utilities with no dependencies on library code. */
+const DEBUG = 3;
+function unzip2(pairs) {
+	const lst1 = [];
+	const lst2 = [];
+	for (const [x, y] of pairs) {
+		lst1.push(x);
+		lst2.push(y);
+	}
+	return [lst1, lst2];
+}
+function zip(xs, ys) {
+	return xs.map((x, i) => [x, ys[i]]);
+}
+function rep(length, value) {
+	if (value instanceof Function) return new Array(length).fill(0).map((_, i) => value(i));
+	return new Array(length).fill(value);
+}
+function prod(arr) {
+	return arr.reduce((acc, x) => acc * x, 1);
+}
+/** Shorthand for integer division, like in Python. */
+function intdiv(a, b) {
+	return Math.floor(a / b);
+}
+/** Clamp `x` to the range `[min, max]`. */
+function clamp(x, min, max) {
+	return Math.max(min, Math.min(max, x));
+}
+/** Check if two objects are deep equal. */
+function deepEqual(a, b) {
+	if (a === b) return true;
+	if (typeof a !== "object" || typeof b !== "object") return false;
+	if (a === null || b === null) return false;
+	if (Object.keys(a).length !== Object.keys(b).length) return false;
+	for (const key of Object.keys(a)) if (!deepEqual(a[key], b[key])) return false;
+	return true;
+}
+/** Splits the list based on a condition, `false` first then `true`. */
+function partitionList(which, array) {
+	const falseList = [];
+	const trueList = [];
+	for (let i = 0; i < which.length; i++) if (which[i]) trueList.push(array[i]);
+	else falseList.push(array[i]);
+	return [falseList, trueList];
+}
+/** Compare two arrays of numbers lexicographically. */
+function lexCompare(a, b) {
+	const minLength = Math.min(a.length, b.length);
+	for (let i = 0; i < minLength; i++) {
+		if (a[i] < b[i]) return -1;
+		if (a[i] > b[i]) return 1;
+	}
+	return a.length - b.length;
+}
+/** Check if an object is a number pair, i.e., a tuple of two numbers. */
+function isNumberPair(x) {
+	return Array.isArray(x) && x.length === 2 && typeof x[0] === "number" && typeof x[1] === "number";
+}
+/** 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`);
+	return axis < 0 ? axis + ndim : axis;
+}
+function range(start, stop, step = 1) {
+	if (stop === void 0) {
+		stop = start;
+		start = 0;
+	}
+	const result = [];
+	for (let i = start; i < stop; i += step) result.push(i);
+	return result;
+}
+function isPermutation(axis, n) {
+	if (axis.length !== n) return false;
+	const seen = /* @__PURE__ */ new Set();
+	for (const x of axis) {
+		if (x < 0 || x >= n) return false;
+		seen.add(x);
+	}
+	return seen.size === n;
+}
+function invertPermutation(axis) {
+	const n = axis.length;
+	if (!isPermutation(axis, n)) throw new Error("invertPermutation: axis is not a permutation");
+	const result = new Array(n);
+	for (let i = 0; i < n; i++) result[axis[i]] = i;
+	return result;
+}
+/** Topologically sort a DAG, given terminal nodes and an ancestor function. */
+function toposort(terminals, parents) {
+	const childCounts = /* @__PURE__ */ new Map();
+	const stack = [...new Set(terminals)];
+	while (true) {
+		const node = stack.pop();
+		if (!node) break;
+		for (const parent of parents(node)) if (childCounts.has(parent)) childCounts.set(parent, childCounts.get(parent) + 1);
+		else {
+			childCounts.set(parent, 1);
+			stack.push(parent);
+		}
+	}
+	for (const node of terminals) childCounts.set(node, childCounts.get(node) - 1);
+	const order = [];
+	const frontier = terminals.filter((n) => !childCounts.get(n));
+	while (true) {
+		const node = frontier.pop();
+		if (!node) break;
+		order.push(node);
+		for (const parent of parents(node)) {
+			const c = childCounts.get(parent) - 1;
+			childCounts.set(parent, c);
+			if (c == 0) frontier.push(parent);
+		}
+	}
+	return order.reverse();
+}
+/**
+* Returns the largest power of 2 less than or equal to `max`.
+*
+* If `hint` is nonzero, it will not return a number greater than the first
+* power of 2 that is greater than or equal to `hint`.
+*/
+function findPow2(hint, max) {
+	if (max < 1) throw new Error("max must be a positive integer");
+	let ret = 1;
+	while (ret < hint && 2 * ret <= max) ret *= 2;
+	return ret;
+}
+function recursiveFlatten(ar) {
+	if (!Array.isArray(ar)) return [ar];
+	return ar.flat(Infinity);
+}
+/** Strip an outermost pair of nested parentheses from an expression, if any. */
+function strip1(str) {
+	if (str[0] === "(" && str[str.length - 1] === ")") return str.slice(1, -1);
+	return str;
+}
+/**
+* Polynomial hashes modulo p are good at avoiding collisions in expectation.
+* Probability-wise, it's good enough to be used for something like
+* deduplicating seen compiler expressions, although it's not adversarial.
+*
+* See https://en.wikipedia.org/wiki/Lagrange%27s_theorem_(number_theory)
+*/
+var FpHash = class FpHash {
+	value = 8773157n;
+	#update(x) {
+		const base = 873192869n;
+		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));
+		else {
+			const ar = new Float64Array([x]);
+			this.#update(new DataView(ar.buffer).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);
+		return this;
+	}
+	static hash(...values) {
+		return new FpHash().update(...values).value;
+	}
+};
+/** Run a function while caching it inline inside a `Map`. */
+function runWithCache(cache, key, thunk) {
+	if (cache.has(key)) return cache.get(key);
+	else {
+		const value = thunk();
+		cache.set(key, value);
+		return value;
+	}
+}
+
+//#endregion
+//#region src/alu.ts
+let DType = /* @__PURE__ */ function(DType$1) {
+	DType$1["Float32"] = "float32";
+	DType$1["Int32"] = "int32";
+	DType$1["Uint32"] = "uint32";
+	DType$1["Bool"] = "bool";
+	DType$1["Complex64"] = "complex64";
+	return DType$1;
+}({});
+const byteWidth = (dtype) => {
+	switch (dtype) {
+		case DType.Float32:
+		case DType.Int32:
+		case DType.Uint32:
+		case DType.Bool: return 4;
+		case DType.Complex64: return 8;
+		default: throw new TypeError(`Unknown dtype: ${dtype}`);
+	}
+};
+const isFloatDtype = (dtype) => dtype === DType.Float32 || dtype === DType.Complex64;
+function dtypedArray(dtype, data) {
+	switch (dtype) {
+		case DType.Float32: return new Float32Array(data);
+		case DType.Int32: return new Int32Array(data);
+		case DType.Uint32: return new Uint32Array(data);
+		case DType.Bool: return new Int32Array(data);
+		default: throw new Error(`Unimplemented dtype: ${dtype}`);
+	}
+}
+/**
+* Mathematical expression on scalar values.
+*
+* This is similiar to and based on tinygrad's UOp class, but it's more specific
+* to just math on scalars. We're doing this to avoid the complexity of a full
+* graph rewrite engine.
+*/
+var AluExp = class AluExp {
+	#hash;
+	#simplified;
+	#range;
+	constructor(op, dtype, src, arg = void 0) {
+		this.op = op;
+		this.dtype = dtype;
+		this.src = src;
+		this.arg = arg;
+		if (AluGroup.RequiredFloat.has(op) && !isFloatDtype(dtype)) throw new TypeError(`Unsupported dtype for ${op}: ${dtype}`);
+		if (op === AluOp.Bitcast && (dtype === DType.Bool || src[0].dtype === DType.Bool || byteWidth(dtype) !== byteWidth(src[0].dtype))) throw new TypeError(`Bitcast from ${src[0].dtype} -> ${dtype}`);
+		if (op === AluOp.Threefry2x32 && (dtype !== DType.Uint32 || src.some((x) => x.dtype !== DType.Uint32))) throw new TypeError("Threefry2x32 requires uint32 types");
+	}
+	static add(a, b) {
+		return new AluExp(AluOp.Add, a.dtype, [a, b]);
+	}
+	static sub(a, b) {
+		return new AluExp(AluOp.Sub, a.dtype, [a, b]);
+	}
+	static mul(a, b) {
+		return new AluExp(AluOp.Mul, a.dtype, [a, b]);
+	}
+	static idiv(a, b) {
+		return new AluExp(AluOp.Idiv, a.dtype, [a, b]);
+	}
+	static mod(a, b) {
+		return new AluExp(AluOp.Mod, a.dtype, [a, b]);
+	}
+	static min(a, b) {
+		return new AluExp(AluOp.Min, a.dtype, [a, b]);
+	}
+	static max(a, b) {
+		return new AluExp(AluOp.Max, a.dtype, [a, b]);
+	}
+	static sin(a) {
+		return new AluExp(AluOp.Sin, a.dtype, [a]);
+	}
+	static cos(a) {
+		return new AluExp(AluOp.Cos, a.dtype, [a]);
+	}
+	static exp(a) {
+		return new AluExp(AluOp.Exp, a.dtype, [a]);
+	}
+	static log(a) {
+		return new AluExp(AluOp.Log, a.dtype, [a]);
+	}
+	static reciprocal(a) {
+		return new AluExp(AluOp.Reciprocal, a.dtype, [a]);
+	}
+	static cast(dtype, a) {
+		if (a.dtype === dtype) return a;
+		return new AluExp(AluOp.Cast, dtype, [a]);
+	}
+	static bitcast(dtype, a) {
+		if (a.dtype === dtype) return a;
+		return new AluExp(AluOp.Bitcast, dtype, [a]);
+	}
+	static threefry2x32(k0, k1, c0, c1, mode = "xor") {
+		return new AluExp(AluOp.Threefry2x32, DType.Uint32, [
+			k0,
+			k1,
+			c0,
+			c1
+		], mode);
+	}
+	static cmplt(a, b) {
+		return new AluExp(AluOp.Cmplt, DType.Bool, [a, b]);
+	}
+	static cmpne(a, b) {
+		return new AluExp(AluOp.Cmpne, DType.Bool, [a, b]);
+	}
+	static where(cond, a, b) {
+		return new AluExp(AluOp.Where, a.dtype, [
+			cond,
+			a,
+			b
+		]);
+	}
+	static const(dtype, value) {
+		if (dtype === DType.Bool) value = Number(Boolean(value));
+		else if (dtype === DType.Int32) value = Math.trunc(value) | 0;
+		else if (dtype === DType.Uint32) value = Math.trunc(value) >>> 0;
+		if (typeof value !== "number") throw new TypeError(`Expected a number for constant, got ${typeof value}: ${value}`);
+		return new AluExp(AluOp.Const, dtype, [], value);
+	}
+	static special(dtype, name, n) {
+		return new AluExp(AluOp.Special, dtype, [], [name, n]);
+	}
+	static variable(dtype, name) {
+		return new AluExp(AluOp.Variable, dtype, [], name);
+	}
+	static globalIndex(dtype, gid, bufidx) {
+		return new AluExp(AluOp.GlobalIndex, dtype, [bufidx], gid);
+	}
+	static globalView(dtype, gid, st, indices) {
+		return new AluExp(AluOp.GlobalView, dtype, indices, [gid, st]);
+	}
+	static i32(value) {
+		return AluExp.const(DType.Int32, value);
+	}
+	static u32(value) {
+		return AluExp.const(DType.Uint32, value);
+	}
+	static f32(value) {
+		return AluExp.const(DType.Float32, value);
+	}
+	static bool(value) {
+		return AluExp.const(DType.Bool, Number(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));
+	}
+	/** Compute a reasonable expression hash with low collision rate. */
+	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);
+		this.#hash = hasher.value;
+		return this.#hash;
+	}
+	hash(state) {
+		state.update(this.getHash());
+	}
+	/** Substitute variables in this AluExp to values. */
+	substitute(variables) {
+		return this.rewrite((exp) => {
+			if (exp.op === AluOp.Variable && Object.hasOwn(variables, exp.arg)) {
+				if (exp.dtype !== variables[exp.arg].dtype) throw new Error(`Type mismatch: ${exp.dtype} vs ${variables[exp.arg].dtype}`);
+				return variables[exp.arg];
+			}
+		});
+	}
+	/** Reindex gid values in this expression as needed. */
+	reindexGids(gidMap) {
+		return this.rewrite((exp) => {
+			if (exp.op === AluOp.GlobalIndex) {
+				const gid = exp.arg;
+				const newGid = gidMap.get(gid);
+				if (newGid !== void 0 && newGid !== gid) return AluExp.globalIndex(exp.dtype, newGid, exp.src[0]);
+			} else if (exp.op === AluOp.GlobalView) {
+				const gid = exp.arg[0];
+				const newGid = gidMap.get(gid);
+				if (newGid !== void 0 && newGid !== gid) return AluExp.globalView(exp.dtype, newGid, exp.arg[1], exp.src);
+			}
+		});
+	}
+	#computeRange() {
+		if (this.#range !== void 0) return this.#range;
+		const src = this.src;
+		const minMax4 = (f) => {
+			const [r1, r2] = [src[0].#computeRange(), src[1].#computeRange()];
+			const values = [
+				f(r1[0], r2[0]),
+				f(r1[0], r2[1]),
+				f(r1[1], r2[0]),
+				f(r1[1], r2[1])
+			];
+			return [Math.min(...values), Math.max(...values)];
+		};
+		let ret;
+		switch (this.op) {
+			case AluOp.Add:
+				ret = [src[0].min + src[1].min, src[0].max + src[1].max];
+				break;
+			case AluOp.Sub:
+				ret = [src[0].min - src[1].max, src[0].max - src[1].min];
+				break;
+			case AluOp.Mul: {
+				ret = minMax4((a, b) => a * b);
+				break;
+			}
+			case AluOp.Idiv: {
+				ret = minMax4((a, b) => Math.floor(a / b));
+				break;
+			}
+			case AluOp.Mod: {
+				let divisorRange = src[1].#computeRange();
+				if (divisorRange[0] <= 0 && divisorRange[1] >= 0) divisorRange = [0, Math.max(-divisorRange[0], divisorRange[1])];
+				const maxDivisor = isFloatDtype(this.dtype) ? divisorRange[1] : divisorRange[1] - 1;
+				ret = [clamp(src[0].min, -maxDivisor, 0), clamp(src[0].max, 0, maxDivisor)];
+				break;
+			}
+			case AluOp.Min:
+				ret = [Math.min(src[0].min, src[1].min), Math.min(src[0].max, src[1].max)];
+				break;
+			case AluOp.Max:
+				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)];
+				break;
+			case AluOp.Cos:
+				ret = [Math.cos(src[0].min), Math.cos(src[0].max)];
+				break;
+			case AluOp.Exp:
+				ret = [Math.exp(src[0].min), Math.exp(src[0].max)];
+				break;
+			case AluOp.Log:
+				ret = [Math.log(src[0].min), Math.log(src[0].max)];
+				break;
+			case AluOp.Reciprocal:
+				if (src[0].min <= 0 && src[0].max >= 0) return [-Infinity, Infinity];
+				ret = [1 / src[0].max, 1 / src[0].min];
+				break;
+			case AluOp.Cast:
+				if (this.dtype === DType.Bool) {
+					const canBeZero = src[0].min <= 0 && src[0].max >= 0;
+					const mustBeZero = src[0].min === 0 && src[0].max === 0;
+					ret = mustBeZero ? [0, 0] : canBeZero ? [0, 1] : [1, 1];
+				} else if (this.dtype === DType.Int32) ret = [Math.trunc(src[0].min), Math.trunc(src[0].max)];
+				else if (this.dtype === DType.Uint32) {
+					const a = Math.trunc(src[0].min);
+					const b = Math.trunc(src[0].max);
+					if (Math.floor(a / 2 ** 32) !== Math.floor(b / 2 ** 32)) ret = [0, -1 >>> 0];
+					else ret = [a % 2 ** 32, b % 2 ** 32];
+				} else ret = [src[0].min, src[0].max];
+				break;
+			case AluOp.Cmplt:
+				ret = [0, 1];
+				break;
+			case AluOp.Cmpne:
+				ret = [0, 1];
+				break;
+			case AluOp.Where:
+				ret = [Math.min(src[1].min, src[2].min), Math.max(src[1].max, src[2].max)];
+				break;
+			case AluOp.Const:
+				ret = [this.arg, this.arg];
+				break;
+			case AluOp.Special:
+				ret = [0, this.arg[1] - 1];
+				break;
+			default: ret = [-Infinity, Infinity];
+		}
+		if (isNaN(ret[0]) || isNaN(ret[1])) ret = [-Infinity, Infinity];
+		if (this.dtype === DType.Bool) {
+			ret[0] = clamp(ret[0], 0, 1);
+			ret[1] = clamp(ret[1], 0, 1);
+		}
+		this.#range = ret;
+		return ret;
+	}
+	get min() {
+		return this.#computeRange()[0];
+	}
+	get max() {
+		return this.#computeRange()[1];
+	}
+	#isConstInt() {
+		return this.op === AluOp.Const && (this.dtype === DType.Int32 || this.dtype === DType.Uint32);
+	}
+	/**
+	* Simplify the expression by replacing any known patterns and deduping
+	* identical subexpressions.
+	*/
+	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 simplified = this.#simplifyInner(cache);
+		const simplifiedHash = simplified.getHash();
+		if (cache.has(simplifiedHash)) {
+			const prevSimplified = cache.get(simplifiedHash);
+			cache.set(hash, prevSimplified);
+			this.#simplified = prevSimplified;
+			return prevSimplified;
+		} else {
+			cache.set(hash, simplified);
+			cache.set(simplifiedHash, simplified);
+			this.#simplified = simplified;
+			return simplified;
+		}
+	}
+	#simplifyInner(cache) {
+		const src = this.src.map((x) => x.simplify(cache));
+		const { op } = this;
+		if (src.every((x) => x.op === AluOp.Const) && !AluGroup.Variable.has(op)) {
+			const newExp$1 = new AluExp(op, this.dtype, src, this.arg);
+			return AluExp.const(this.dtype, newExp$1.evaluate({}));
+		}
+		if (op !== AluOp.Const && this.min === this.max) return AluExp.const(this.dtype, this.min);
+		if (AluGroup.Binary.has(op)) for (let i = 0; i < 2; i++) {
+			if (src[i].op !== AluOp.Const) continue;
+			const x = src[i].arg;
+			if (op === AluOp.Add && x === 0) return src[1 - i];
+			if (op === AluOp.Sub && i === 1 && x === 0) return src[1 - i];
+			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.Add || op === AluOp.Sub) && src[1].op === AluOp.Mul) {
+			const [a, b] = src[1].src;
+			const opNeg = op === AluOp.Add ? AluOp.Sub : AluOp.Add;
+			if (a.op === AluOp.Const && a.arg === -1) return new AluExp(opNeg, this.dtype, [src[0], b]);
+			else if (b.op === AluOp.Const && b.arg === -1) return new AluExp(opNeg, this.dtype, [src[0], a]);
+		}
+		if (op === AluOp.Mod && src[1].op === AluOp.Const && src[0].min >= 0 && src[0].max < src[1].arg) return src[0];
+		if (op === AluOp.Add && src[0].op === AluOp.Mul && src[0].src[1].#isConstInt() && src[1].op === AluOp.Mod && src[1].src[1].#isConstInt() && src[0].src[1].arg === src[1].src[1].arg) {
+			const [mul, mod] = src;
+			const check = (exp) => {
+				return exp.op === AluOp.Idiv && exp.src[1].#isConstInt() && exp.src[1].arg === mod.src[1].arg && exp.src[0] === mod.src[0];
+			};
+			if (check(mul.src[0])) return mod.src[0];
+			if (mul.src[0].op === AluOp.Mod) {
+				const [x, y] = mul.src[0].src;
+				if (check(x)) return AluExp.mod(mod.src[0], AluExp.mul(mod.src[1], y)).simplify(cache);
+			}
+		}
+		if (op === AluOp.Idiv && src[1].#isConstInt()) {
+			const [numer, denom] = src;
+			const B = denom.arg;
+			for (let i = 0; i < 2; i++) {
+				if (numer.op === AluOp.Mul && numer.src[i].#isConstInt()) {
+					const A = numer.src[i].arg;
+					if (A % B === 0) {
+						let ret = numer.src[1 - i];
+						if (A / B !== 1) ret = AluExp.mul(ret, AluExp.i32(A / B));
+						return ret.simplify(cache);
+					}
+				}
+				for (let j = 0; j < 2; j++) if (numer.op === AluOp.Add && numer.src[j].op === AluOp.Mul && numer.src[j].src[i].#isConstInt()) {
+					const A = numer.src[j].src[i].arg;
+					if (A % B === 0) {
+						let ret = numer.src[j].src[1 - i];
+						if (A / B !== 1) ret = AluExp.mul(ret, AluExp.i32(A / B));
+						ret = AluExp.add(ret, AluExp.idiv(numer.src[1 - j], B));
+						return ret.simplify(cache);
+					}
+				}
+			}
+		}
+		if (op === AluOp.Mod && src[1].#isConstInt() && src[1].arg > 0 && src[0].min >= 0) {
+			const [numer, denom] = src;
+			const B = denom.arg;
+			for (let i = 0; i < 2; i++) if (numer.op === AluOp.Add && numer.src[i].#isConstInt()) {
+				const A = numer.src[i].arg;
+				let ret = numer.src[1 - i];
+				if (A % B !== 0) ret = AluExp.add(ret, AluExp.i32(A % B));
+				return ret.simplify(cache);
+			}
+		}
+		if (op === AluOp.Cmplt) {
+			if (src[0].min >= src[1].max) return AluExp.const(DType.Bool, false);
+			if (src[0].max < src[1].min) return AluExp.const(DType.Bool, true);
+		}
+		if (op === AluOp.Cmpne) {
+			if (src[0].max < src[1].min || src[0].min > src[1].max) return AluExp.const(DType.Bool, true);
+		}
+		if (op === AluOp.Where) {
+			if (src[0].max === 0) return src[2];
+			if (src[0].min === 1) return src[1];
+		}
+		const newExp = src.every((s, i) => s === this.src[i]) ? this : new AluExp(op, this.dtype, src, this.arg);
+		return newExp;
+	}
+	/** Resolve this to a value, or `undefined` if not possible. */
+	resolve() {
+		const x = this.simplify();
+		if (x.op === AluOp.Const) return x.arg;
+		return void 0;
+	}
+	/**
+	* Evaluate the expression on CPU, returning the result.
+	*
+	* Typically you would compile the AluExp as a representation to a lower-level
+	* language. This is just to define the semantics and help debug.
+	*
+	* Note that the representation of Bool is as a number (0 or 1) here.
+	*/
+	evaluate(context, globals) {
+		if (AluGroup.Binary.has(this.op) || AluGroup.Compare.has(this.op)) {
+			const x = this.src[0].evaluate(context, globals);
+			const y = this.src[1].evaluate(context, globals);
+			switch (this.op) {
+				case AluOp.Add: return this.dtype === DType.Bool ? Number(x || y) : x + y;
+				case AluOp.Sub: return x - y;
+				case AluOp.Mul: return this.dtype === DType.Bool ? Number(x && y) : x * y;
+				case AluOp.Idiv: return Math.trunc(x / y);
+				case AluOp.Mod: return x % y;
+				case AluOp.Min: return Math.min(x, y);
+				case AluOp.Max: return Math.max(x, y);
+				case AluOp.Cmplt: return Number(x < y);
+				case AluOp.Cmpne: return Number(x != y);
+				default: throw new Error(`Missing implemementation for ${this.op}`);
+			}
+		}
+		if (AluGroup.Unary.has(this.op)) {
+			const x = this.src[0].evaluate(context, globals);
+			switch (this.op) {
+				case AluOp.Sin: return Math.sin(x);
+				case AluOp.Cos: return Math.cos(x);
+				case AluOp.Exp: return Math.exp(x);
+				case AluOp.Log: return Math.log(x);
+				case AluOp.Reciprocal: return 1 / x;
+				case AluOp.Cast: if (this.dtype === DType.Int32) return Math.trunc(x) | 0;
+				else if (this.dtype === DType.Uint32) return Math.trunc(x) >>> 0;
+				else if (this.dtype === DType.Float32) return x;
+				else if (this.dtype === DType.Bool) return Number(Boolean(x));
+				else throw new Error(`Unsupported cast to ${this.dtype}`);
+				case AluOp.Bitcast: {
+					const buf = new ArrayBuffer(byteWidth(this.dtype));
+					const view = new DataView(buf);
+					const fromType = this.src[0].dtype;
+					if (fromType === DType.Float32) view.setFloat32(0, x, true);
+					else if (fromType === DType.Int32) view.setInt32(0, x, true);
+					else if (fromType === DType.Uint32) view.setUint32(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 throw new Error(`Unsupported bitcast to ${this.dtype}`);
+				}
+				default: throw new Error(`Missing implemementation for ${this.op}`);
+			}
+		}
+		switch (this.op) {
+			case AluOp.Where: return this.src[0].evaluate(context, globals) ? this.src[1].evaluate(context, globals) : this.src[2].evaluate(context, globals);
+			case AluOp.Threefry2x32: {
+				const [k0, k1, c0, c1] = this.src.map((x) => x.evaluate(context, globals));
+				const [x0, x1] = threefry2x32(k0, k1, c0, c1);
+				if (this.arg === "xor") return (x0 ^ x1) >>> 0;
+				else if (this.arg === 0) return x0;
+				else if (this.arg === 1) return x1;
+				else throw new Error(`Invalid Threefry2x32 mode: ${this.arg}`);
+			}
+			case AluOp.Const: return this.arg;
+			case AluOp.Special: {
+				const x = context[this.arg[0]];
+				if (x === void 0) throw new Error(`Missing special: ${this.arg[0]}`);
+				return x;
+			}
+			case AluOp.Variable: {
+				const x = context[this.arg];
+				if (x === void 0) throw new Error(`Missing variable: ${this.arg}`);
+				return x;
+			}
+			case AluOp.GlobalIndex: {
+				if (!globals) throw new Error("Missing globals function");
+				const gid = this.arg;
+				const bufidx = this.src[0].evaluate(context, globals);
+				return globals(gid, bufidx);
+			}
+			case AluOp.GlobalView: {
+				if (!globals) throw new Error("Missing globals function");
+				const gid = this.arg[0];
+				const st = this.arg[1];
+				const [iexpr, vexpr] = st.toAluExp(this.src);
+				if (vexpr.evaluate(context, globals)) {
+					const bufidx = iexpr.evaluate(context, globals);
+					return globals(gid, bufidx);
+				} else return 0;
+			}
+			default: throw new Error(`Missing implemementation for ${this.op}`);
+		}
+	}
+	/** Get this expression in debug format as a string. */
+	toString() {
+		const BIN_SYM = {
+			[AluOp.Add]: "+",
+			[AluOp.Sub]: "-",
+			[AluOp.Mul]: "*",
+			[AluOp.Idiv]: "/",
+			[AluOp.Mod]: "%"
+		};
+		const CMP_SYM = {
+			[AluOp.Cmplt]: "<",
+			[AluOp.Cmpne]: "!="
+		};
+		const UNARY_SYM = {
+			[AluOp.Sin]: "sin",
+			[AluOp.Cos]: "cos",
+			[AluOp.Exp]: "exp",
+			[AluOp.Log]: "log",
+			[AluOp.Reciprocal]: "1/"
+		};
+		return this.fold((node, parts) => {
+			switch (node.op) {
+				case AluOp.Const: return "" + (node.dtype === DType.Bool ? Boolean(node.arg) : node.arg);
+				case AluOp.Variable: return `$${node.arg}:${node.dtype}`;
+				case AluOp.Special: {
+					const [name, n] = node.arg;
+					return `#${name}{${n}}`;
+				}
+				case AluOp.GlobalIndex: return `G_${node.arg}<${node.dtype}>[${strip1(parts[0])}]`;
+				case AluOp.GlobalView: {
+					const [gid, st] = node.arg;
+					const shape = st.shape.join(",");
+					const lastStrides = st.lastStrides.join(",");
+					const cont = st.contiguous ? "c" : "nc";
+					return `GV_${gid}<${node.dtype}>{${shape}:${lastStrides}:${cont}}[${parts.map(strip1).join(", ")}]`;
+				}
+			}
+			if (BIN_SYM[node.op]) return `(${parts[0]} ${BIN_SYM[node.op]} ${parts[1]})`;
+			if (CMP_SYM[node.op]) return `(${parts[0]} ${CMP_SYM[node.op]} ${parts[1]})`;
+			if (UNARY_SYM[node.op]) return `${UNARY_SYM[node.op]}${parts[0]}`;
+			if (node.op === AluOp.Cast) return `Cast<${node.dtype}>(${strip1(parts[0])})`;
+			if (node.op === AluOp.Bitcast) return `Bitcast<${node.dtype}>(${strip1(parts[0])})`;
+			return `${node.op}(${parts.map(strip1).join(", ")})`;
+		});
+	}
+	/** Generic fold() operation with a reducer over the expression tree. */
+	fold(reducer) {
+		const visited = /* @__PURE__ */ new Map();
+		const recurse = (exp) => {
+			if (visited.has(exp)) return visited.get(exp);
+			const mappedSrc = exp.src.map((s) => recurse(s));
+			const result = reducer(exp, mappedSrc);
+			visited.set(exp, result);
+			return result;
+		};
+		return recurse(this);
+	}
+	/** Rewrite the expression recursively using a visitor. */
+	rewrite(visitor) {
+		return this.fold((exp, newSrc) => {
+			if (newSrc.length === exp.src.length && newSrc.every((s, i) => s === exp.src[i])) return visitor(exp) ?? exp;
+			else {
+				const newExp = new AluExp(exp.op, exp.dtype, newSrc, exp.arg);
+				return visitor(newExp) ?? newExp;
+			}
+		});
+	}
+	/** Collect all nodes that satisfy a predicate. */
+	collect(predicate) {
+		const result = [];
+		this.fold((exp) => {
+			if (predicate(exp)) result.push(exp);
+		});
+		return result;
+	}
+};
+/** Symbolic form for each mathematical operation. */
+let AluOp = /* @__PURE__ */ function(AluOp$1) {
+	AluOp$1["Add"] = "Add";
+	AluOp$1["Sub"] = "Sub";
+	AluOp$1["Mul"] = "Mul";
+	AluOp$1["Idiv"] = "Idiv";
+	AluOp$1["Mod"] = "Mod";
+	AluOp$1["Min"] = "Min";
+	AluOp$1["Max"] = "Max";
+	AluOp$1["Sin"] = "Sin";
+	AluOp$1["Cos"] = "Cos";
+	AluOp$1["Exp"] = "Exp";
+	AluOp$1["Log"] = "Log";
+	AluOp$1["Reciprocal"] = "Reciprocal";
+	AluOp$1["Cast"] = "Cast";
+	AluOp$1["Bitcast"] = "Bitcast";
+	AluOp$1["Cmplt"] = "Cmplt";
+	AluOp$1["Cmpne"] = "Cmpne";
+	AluOp$1["Where"] = "Where";
+	AluOp$1["Threefry2x32"] = "Threefry2x32";
+	AluOp$1["Const"] = "Const";
+	AluOp$1["Special"] = "Special";
+	AluOp$1["Variable"] = "Variable";
+	AluOp$1["GlobalIndex"] = "GlobalIndex";
+	AluOp$1["GlobalView"] = "GlobalView";
+	return AluOp$1;
+}({});
+const AluGroup = {
+	Binary: new Set([
+		AluOp.Add,
+		AluOp.Sub,
+		AluOp.Mul,
+		AluOp.Idiv,
+		AluOp.Mod,
+		AluOp.Min,
+		AluOp.Max
+	]),
+	Unary: new Set([
+		AluOp.Sin,
+		AluOp.Cos,
+		AluOp.Exp,
+		AluOp.Log,
+		AluOp.Reciprocal,
+		AluOp.Cast,
+		AluOp.Bitcast
+	]),
+	Compare: new Set([AluOp.Cmplt, AluOp.Cmpne]),
+	Variable: new Set([
+		AluOp.Special,
+		AluOp.Variable,
+		AluOp.GlobalIndex,
+		AluOp.GlobalView
+	]),
+	Reduce: new Set([
+		AluOp.Add,
+		AluOp.Mul,
+		AluOp.Min,
+		AluOp.Max
+	]),
+	RequiredFloat: new Set([
+		AluOp.Sin,
+		AluOp.Cos,
+		AluOp.Exp,
+		AluOp.Log,
+		AluOp.Reciprocal
+	])
+};
+/** Common variables that can be substituted in expressions. */
+const AluVar = {
+	gidx: AluExp.variable(DType.Int32, "gidx"),
+	ridx: AluExp.variable(DType.Int32, "ridx"),
+	acc: (dtype) => AluExp.variable(dtype, "acc"),
+	idx: AluExp.variable(DType.Int32, "idx"),
+	unroll: AluExp.variable(DType.Int32, "unroll"),
+	upcast: AluExp.variable(DType.Int32, "upcast")
+};
+/**
+* Description of a kernel to be compiled.
+*
+* Each of these can be processed by a backend into some lower-level
+* representation. It consists of one or more fused operations, optionally
+* indexing into a buffer.
+*/
+var Kernel = class {
+	constructor(nargs, size, exp, reduction) {
+		this.nargs = nargs;
+		this.size = size;
+		this.exp = exp;
+		this.reduction = reduction;
+		this.exp = exp.simplify();
+	}
+	hash(state) {
+		state.update(this.nargs, this.size, this.exp, this.reduction);
+	}
+	pprint() {
+		let details = PPrint.pp(`exp = ${this.exp}`);
+		details = details.concat(PPrint.pp(`size = ${this.size}`));
+		if (this.reduction) details = details.concat(PPrint.pp(`reduction = ${this.reduction}`));
+		return PPrint.pp("{ ").stack(details).stack(PPrint.pp(" }"));
+	}
+	toString() {
+		return this.pprint().toString();
+	}
+	/** The dtype of the values output by this kernel. */
+	get dtype() {
+		if (this.reduction) return this.reduction.fusion.dtype;
+		else return this.exp.dtype;
+	}
+	/** The number of bytes in the output array when evaluating this kernel. */
+	get bytes() {
+		return this.size * byteWidth(this.dtype);
+	}
+};
+/**
+* Description of a reduction.
+*
+* The strategy of jax-js backends is to either handle a standard operation that
+* is dispatched in a vectorized way over an array, or to reduce over one axis
+* of some computation. This is a description of the reduction.
+*
+* Reduction only supports a few operations, and only over one axis. Users can
+* always `flatten()` the array before reducing if needed.
+*
+* The backend is responsible for implementing the reduction in a way that
+* minimizes the number of global memory loads, for efficiency. This involves
+* passing through some optimization strategy. But optimizations are not coded
+* at this level since they depend on GPU, versus CPU or Wasm.
+*/
+var Reduction = class {
+	constructor(dtype, op, size, fusion = AluVar.acc(dtype)) {
+		this.dtype = dtype;
+		this.op = op;
+		this.size = size;
+		this.fusion = fusion;
+		if (!AluGroup.Reduce.has(op)) throw new TypeError(`Unsupported reduction: ${op}`);
+	}
+	hash(state) {
+		state.update(this.dtype, this.op, this.size, this.fusion);
+	}
+	toString() {
+		return `${this.op}{${this.size}} -> ${this.fusion}`;
+	}
+	/** Get the identity for this reduction operation. */
+	get identity() {
+		if (this.dtype === DType.Bool) return this.op === AluOp.Add || this.op === AluOp.Max ? false : true;
+		else if (this.dtype === DType.Int32) {
+			if (this.op === AluOp.Add) return 0;
+			else if (this.op === AluOp.Mul) return 1;
+			else if (this.op === AluOp.Min) return -1 >>> 1;
+			else if (this.op === AluOp.Max) return 1 << 31;
+		} else if (this.dtype === DType.Uint32) {
+			if (this.op === AluOp.Add) return 0;
+			else if (this.op === AluOp.Mul) return 1;
+			else if (this.op === AluOp.Min) return -1 >>> 0;
+			else if (this.op === AluOp.Max) return 0;
+		} else if (this.dtype === DType.Float32) {
+			if (this.op === AluOp.Add) return 0;
+			else if (this.op === AluOp.Mul) return 1;
+			else if (this.op === AluOp.Min) return Infinity;
+			else if (this.op === AluOp.Max) return -Infinity;
+		}
+		throw new TypeError(`Unsupported reduction: ${this.op} ${this.dtype}`);
+	}
+	/** 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);
+			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);
+			else if (this.op === AluOp.Mul) return values.reduce((a, b) => a * b | 0, 1);
+			else if (this.op === AluOp.Min) return values.reduce((a, b) => Math.min(a, b), -1 >>> 1);
+			else if (this.op === AluOp.Max) return values.reduce((a, b) => Math.max(a, b), 1 << 31);
+		} else if (this.dtype === DType.Uint32) {
+			if (this.op === AluOp.Add) return values.reduce((a, b) => a + b >>> 0, 0);
+			else if (this.op === AluOp.Mul) return values.reduce((a, b) => a * b >>> 0, 1);
+			else if (this.op === AluOp.Min) return values.reduce((a, b) => Math.min(a, b), -1 >>> 0);
+			else if (this.op === AluOp.Max) return values.reduce((a, b) => Math.max(a, b), 0);
+		} else if (this.dtype === DType.Float32) {
+			if (this.op === AluOp.Add) return values.reduce((a, b) => a + b, 0);
+			else if (this.op === AluOp.Mul) return values.reduce((a, b) => a * b, 1);
+			else if (this.op === AluOp.Min) return values.reduce((a, b) => Math.min(a, b), Infinity);
+			else if (this.op === AluOp.Max) return values.reduce((a, b) => Math.max(a, b), -Infinity);
+		}
+		throw new TypeError(`Unsupported reduction: ${this.op} ${this.dtype}`);
+	}
+};
+/** Expression for accessing `indices` in input array with the given shape. */
+function accessorGlobal(dtype, gid, st, indices) {
+	const [index, valid] = st.toAluExp(indices);
+	return AluExp.where(valid, AluExp.globalIndex(dtype, gid, index), AluExp.const(dtype, 0));
+}
+/** Expression for accessing `indices` in an array recipe with variable "idx". */
+function accessorAluExp(dtype, exp, st, indices) {
+	const [index, valid] = st.toAluExp(indices);
+	return AluExp.where(valid, exp.substitute({ idx: index }), AluExp.const(dtype, 0));
+}
+function threefry2x32(k0, k1, c0, c1) {
+	const rotl32 = (x, r) => (x << r | x >>> 32 - r) >>> 0;
+	const ks0 = k0 >>> 0;
+	const ks1 = k1 >>> 0;
+	const ks2 = (ks0 ^ ks1 ^ 466688986) >>> 0;
+	let x0 = c0 + ks0 >>> 0;
+	let x1 = c1 + ks1 >>> 0;
+	x0 = x0 + x1 >>> 0, x1 = rotl32(x1, 13) ^ x0;
+	x0 = x0 + x1 >>> 0, x1 = rotl32(x1, 15) ^ x0;
+	x0 = x0 + x1 >>> 0, x1 = rotl32(x1, 26) ^ x0;
+	x0 = x0 + x1 >>> 0, x1 = rotl32(x1, 6) ^ x0;
+	x0 = x0 + ks1 >>> 0;
+	x1 = x1 + ks2 + 1 >>> 0;
+	x0 = x0 + x1 >>> 0, x1 = rotl32(x1, 17) ^ x0;
+	x0 = x0 + x1 >>> 0, x1 = rotl32(x1, 29) ^ x0;
+	x0 = x0 + x1 >>> 0, x1 = rotl32(x1, 16) ^ x0;
+	x0 = x0 + x1 >>> 0, x1 = rotl32(x1, 24) ^ x0;
+	x0 = x0 + ks2 >>> 0;
+	x1 = x1 + ks0 + 2 >>> 0;
+	x0 = x0 + x1 >>> 0, x1 = rotl32(x1, 13) ^ x0;
+	x0 = x0 + x1 >>> 0, x1 = rotl32(x1, 15) ^ x0;
+	x0 = x0 + x1 >>> 0, x1 = rotl32(x1, 26) ^ x0;
+	x0 = x0 + x1 >>> 0, x1 = rotl32(x1, 6) ^ x0;
+	x0 = x0 + ks0 >>> 0;
+	x1 = x1 + ks1 + 3 >>> 0;
+	x0 = x0 + x1 >>> 0, x1 = rotl32(x1, 17) ^ x0;
+	x0 = x0 + x1 >>> 0, x1 = rotl32(x1, 29) ^ x0;
+	x0 = x0 + x1 >>> 0, x1 = rotl32(x1, 16) ^ x0;
+	x0 = x0 + x1 >>> 0, x1 = rotl32(x1, 24) ^ x0;
+	x0 = x0 + ks1 >>> 0;
+	x1 = x1 + ks2 + 4 >>> 0;
+	x0 = x0 + x1 >>> 0, x1 = rotl32(x1, 13) ^ x0;
+	x0 = x0 + x1 >>> 0, x1 = rotl32(x1, 15) ^ x0;
+	x0 = x0 + x1 >>> 0, x1 = rotl32(x1, 26) ^ x0;
+	x0 = x0 + x1 >>> 0, x1 = rotl32(x1, 6) ^ x0;
+	x0 = x0 + ks2 >>> 0;
+	x1 = x1 + ks0 + 5 >>> 0;
+	return [x0, x1];
+}
+
+//#endregion
+//#region src/shape.ts
+const jstr = JSON.stringify;
+/** Remove "1" dimensions from the strides list. */
+function canonicalizeStrides(shape, strides) {
+	const newStrides = [];
+	for (let i = 0; i < shape.length; i++) if (shape[i] === 1) newStrides.push(0);
+	else newStrides.push(strides[i]);
+	return newStrides;
+}
+/** Get the strides for a shape in default row-major order. */
+function defaultStrides(shape) {
+	if (shape.length === 0) return [];
+	const strides = rep(shape.length, 1);
+	for (let i = shape.length - 1; i > 0; i--) strides[i - 1] = shape[i] * strides[i];
+	return canonicalizeStrides(shape, strides);
+}
+/** Merge contiguous subparts or zero-strided dimensions in a view. */
+function mergeDims(shape, strides, mask) {
+	if (shape.length === 0) return [];
+	if (shape.length !== strides.length || mask && shape.length !== mask.length) throw new Error("internal: invalid args to mergeDims");
+	const ret = [[
+		shape[0],
+		strides[0],
+		strides[0] !== 0 ? shape[0] : 0
+	]];
+	let merging = mask ? mask[0][1] - mask[0][0] === 1 : shape[0] === 1;
+	for (let i = 1; i < shape.length; i++) {
+		const [s, st] = [shape[i], strides[i]];
+		if (s === 1) continue;
+		const [lastS, lastSt, lastPreExpandS] = ret[ret.length - 1];
+		if (merging || lastSt === s * st) ret[ret.length - 1] = [
+			lastS * s,
+			st,
+			merging ? s : lastPreExpandS * s
+		];
+		else ret.push([
+			s,
+			st,
+			s
+		]);
+		merging = mask ? mask[i][1] - mask[i][0] === 1 : false;
+	}
+	return ret;
+}
+/** Return the new mask if a reshape if possible, otherwise `null`. */
+function reshapeMask(maskInput, oldShape, newShape) {
+	const newMask = [];
+	let rMasksI = maskInput.length;
+	let rShapeI = oldShape.length;
+	let rNewShapeI = newShape.length;
+	const rMasks = () => rMasksI ? maskInput[--rMasksI] : [0, 1];
+	const rShape = () => rShapeI ? oldShape[--rShapeI] : 1;
+	const rNewShape = () => rNewShapeI ? newShape[--rNewShapeI] : 1;
+	let currStride = 1;
+	let [oldDim, newDim, mask] = [
+		rShape(),
+		rNewShape(),
+		rMasks()
+	];
+	while (newMask.length < newShape.length) {
+		const [l, r] = mask;
+		const nextStride = newDim * currStride;
+		if (oldDim === nextStride) {
+			newMask.push([intdiv(l, currStride), intdiv(r - 1, currStride) + 1]);
+			currStride = 1;
+			[oldDim, newDim, mask] = [
+				rShape(),
+				rNewShape(),
+				rMasks()
+			];
+		} else if (oldDim > nextStride) {
+			if (oldDim % nextStride !== 0) return null;
+			if ((l % nextStride !== 0 || r % nextStride !== 0) && intdiv(l, nextStride) !== intdiv(r - 1, nextStride)) return null;
+			newMask.push([intdiv(l % nextStride, currStride), intdiv((r - 1) % nextStride, currStride) + 1]);
+			[currStride, newDim] = [nextStride, rNewShape()];
+		} else {
+			const nextMask = rMasks();
+			if (!deepEqual(mask, [0, oldDim]) && l !== r && nextMask[1] - nextMask[0] !== 1) return null;
+			mask = [nextMask[0] * oldDim + l, (nextMask[1] - 1) * oldDim + r];
+			oldDim *= rShape();
+		}
+	}
+	return newMask.reverse();
+}
+/**
+* A multidimensional view into memory. An array can be thought of as the
+* combination of a linear buffer of memory, along with a `View`.
+*
+* Formula for getting a data point is basically:
+*   1. Check if ∀i. 0 <= dim[i] < shape[i], otherwise out of bounds.
+*   2. If mask exists, and ∃i. dim[i] ∉ mask[i], return 0.
+*   2. Otherwise, look at this memory address: offset + ∑(strides[i] * dim[i]).
+*/
+var View = class View {
+	#size;
+	#contiguous;
+	constructor(shape, strides, offset, mask) {
+		this.shape = shape;
+		this.strides = strides;
+		this.offset = offset;
+		this.mask = mask;
+	}
+	static create(shape, strides, offset = 0, mask = null) {
+		if (shape.some((s) => s < 0)) throw new Error("View shape must be non-negative");
+		strides = strides ? canonicalizeStrides(shape, strides) : defaultStrides(shape);
+		if (shape.includes(0)) return new View(shape, rep(shape.length, 0), 0, null);
+		if (mask !== null && mask.every(([b, e], i) => b === 0 && e === shape[i])) mask = null;
+		if (mask !== null) {
+			const elimDims = [];
+			let hasNoData = false;
+			for (let i = 0; i < shape.length; i++) {
+				const [b, e] = mask[i];
+				if (b + 1 >= e) elimDims.push(i);
+				if (b >= e) hasNoData = true;
+			}
+			if (elimDims.length) {
+				if (hasNoData) {
+					strides = rep(shape.length, 0);
+					offset = 0;
+					mask = rep(shape.length, () => [0, 0]);
+				}
+				for (const i of elimDims) {
+					offset += strides[i] * mask[i][0];
+					strides[i] = 0;
+				}
+			}
+		}
+		return new View(shape, strides, offset, mask);
+	}
+	get ndim() {
+		return this.shape.length;
+	}
+	get size() {
+		if (this.#size === void 0) this.#size = prod(this.shape);
+		return this.#size;
+	}
+	/** Whether this is a default, contiguous, unaltered view of the data (identity). */
+	get contiguous() {
+		if (this.#contiguous === void 0) this.#contiguous = this.size === 0 || this.offset === 0 && this.mask === null && deepEqual(this.strides, defaultStrides(this.shape));
+		return this.#contiguous;
+	}
+	/** Produce an AluExp for evaluating this view at an index. */
+	toAluExp(idxs) {
+		let iexpr = AluExp.i32(this.offset);
+		let vexpr = AluExp.bool(true);
+		for (let i = this.ndim - 1; i >= 0; i--) {
+			const idx = idxs[i];
+			if (this.shape[i] !== 1 && this.strides[i] !== 0) iexpr = AluExp.add(AluExp.mul(idx, AluExp.i32(this.strides[i])), iexpr);
+			if (this.mask) {
+				if (this.mask[i][0] !== 0) vexpr = AluExp.mul(AluExp.cmplt(idx, AluExp.i32(this.mask[i][0])).not(), vexpr);
+				if (this.mask[i][1] !== this.shape[i]) vexpr = AluExp.mul(AluExp.cmplt(idx, AluExp.i32(this.mask[i][1])), vexpr);
+			}
+		}
+		return [iexpr, vexpr];
+	}
+	/**
+	* Try to compose this view with another one. `this` view is applied first,
+	* followed by the argument. If this is not possible for the specific views,
+	* return `null` instead.
+	*
+	* If composable, return a combined view with the same shape as `v1`.
+	*
+	* This is very tricky. The shapes of v1 and v2 may be different, and in that
+	* case, we do some math to figure out whether they're compatible.
+	*/
+	compose(v1) {
+		const v2 = this;
+		if (v2.contiguous) return v1;
+		if (v1.contiguous) {
+			if (deepEqual(v1.shape, v2.shape)) return v2;
+			if (v1.size === v2.size) {
+				const ret = v2.reshape(v1.shape);
+				if (ret !== null) return ret;
+			}
+		}
+		if (v1.mask !== null) {
+			const newV1 = v1.shrink(v1.mask);
+			const merged = v2.compose(newV1);
+			return merged ? merged.pad(zip(v1.mask, v1.shape).map(([m, s]) => [m[0], s - m[1]])) : null;
+		}
+		const origin = unravel(v2.shape, v1.offset);
+		const terms = rep(v2.ndim, () => []);
+		const strides = rep(v1.ndim, 0);
+		for (let d1 = 0; d1 < v1.strides.length; d1++) {
+			const st = v1.strides[d1];
+			if (st === 0) continue;
+			const unravelOffset = unravel(v2.shape, v1.offset + st);
+			for (let d2 = 0; d2 < v2.ndim; d2++) {
+				const o = origin[d2];
+				const diff = unravelOffset[d2] - o;
+				if (diff === 0) continue;
+				terms[d2].push([d1, diff]);
+				strides[d1] += diff * v2.strides[d2];
+			}
+		}
+		let [mergedSize, mergedTermMin, mergedTermMax] = [
+			1,
+			0,
+			0
+		];
+		const extents = [];
+		for (let i = v2.ndim - 1; i >= 0; i--) {
+			const term = terms[i];
+			const s = v2.shape[i];
+			let [tmin, tmax] = [origin[i], origin[i]];
+			for (const [d1, s1] of term) if (s1 > 0) tmax += (v1.shape[d1] - 1) * s1;
+			else if (s1 < 0) tmin += (v1.shape[d1] - 1) * s1;
+			mergedTermMin += tmin * mergedSize;
+			mergedTermMax += tmax * mergedSize;
+			mergedSize *= s;
+			if (mergedTermMin >= 0 && mergedTermMax < mergedSize) {
+				extents.push([
+					mergedSize,
+					mergedTermMin,
+					mergedTermMax
+				]);
+				[mergedSize, mergedTermMin, mergedTermMax] = [
+					1,
+					0,
+					0
+				];
+			}
+		}
+		if (mergedTermMin !== 0 || mergedTermMax !== 0) return null;
+		extents.reverse();
+		const v2Shape = extents.map(([s]) => s);
+		if (!deepEqual(v2Shape, v2.shape)) {
+			const reshapedV2 = v2.reshape(v2Shape);
+			if (reshapedV2 === null) return null;
+			if (!deepEqual(reshapedV2.shape, v2.shape)) return reshapedV2.compose(v1);
+		}
+		if (v2.mask !== null) {
+			const newB = rep(v1.ndim, 0);
+			const newE = v1.shape.slice();
+			let bad = false;
+			for (let d2 = 0; d2 < v2.ndim; d2++) {
+				const [b, e] = v2.mask[d2];
+				const o = origin[d2];
+				const term = terms[d2];
+				const [_, tmin, tmax] = extents[d2];
+				if (b <= tmin && tmax < e) continue;
+				if (term.length !== 1) if (term.length === 0 && newE.length) newE[0] = 0;
+				else bad = true;
+				else {
+					const [d1, s1] = term[0];
+					newB[d1] = Math.max(newB[d1], Math.ceil((s1 > 0 ? b - o : e - o - 1) / s1));
+					newE[d1] = Math.min(newE[d1], Math.floor((s1 < 0 ? b - o : e - o - 1) / s1) + 1);
+				}
+			}
+			for (let d1 = 0; d1 < v1.ndim; d1++) if (newB[d1] !== 0 || newE[d1] !== v1.shape[d1]) return v2.compose(View.create(v1.shape, v1.strides, v1.offset, zip(newB, newE)));
+			if (bad) return null;
+		}
+		let finalOffset = v2.offset;
+		for (let d2 = 0; d2 < v2.ndim; d2++) finalOffset += origin[d2] * v2.strides[d2];
+		return View.create(v1.shape, strides, finalOffset, null);
+	}
+	/** Attempt to simplify this view into a smaller reshaped form. */
+	minify() {
+		const minShape = mergeDims(this.shape, this.strides, this.mask).map((x) => x[0]);
+		const nv = this.reshape(minShape);
+		return nv ? nv : this;
+	}
+	/** Pad the view with zeros on each dimension. */
+	pad(arg) {
+		if (arg.length !== this.ndim || !arg.every(([b, e]) => b >= 0 && e >= 0)) throw new Error(`invalid pad ${jstr(arg)} for ${jstr(this.shape)}`);
+		if (arg.every(([b, e]) => b === 0 && e === 0)) return this;
+		const zvarg = arg.map(([b, e], i) => [-b, this.shape[i] + e]);
+		const mask = arg.map(([b, _e], i) => [b, this.shape[i] + b]);
+		return this.#unsafeResize(zvarg, mask);
+	}
+	/** Shrink the view by taking a subarray. */
+	shrink(arg) {
+		if (arg.length !== this.ndim || !arg.every(([b, e], i) => 0 <= b && b <= e && e <= this.shape[i])) throw new Error(`invalid shrink ${jstr(arg)} for ${jstr(this.shape)}`);
+		return this.#unsafeResize(arg);
+	}
+	#unsafeResize(arg, mask) {
+		const offset = this.strides.map((s, i) => s * arg[i][0]).reduce((a, b) => a + b, 0);
+		if (this.mask) {
+			const nmask = this.mask.map(([mx, my], i) => [Math.max(0, Math.min(mx - arg[i][0], arg[i][1] - arg[i][0])), Math.max(0, Math.min(my - arg[i][0], arg[i][1] - arg[i][0]))]);
+			mask = mask ? mask.map(([mx, my], i) => [Math.max(mx, nmask[i][0]), Math.min(my, nmask[i][1])]) : nmask;
+		}
+		return View.create(arg.map(([b, e]) => e - b), this.strides, this.offset + offset, mask);
+	}
+	/** Expand one or more axes with length "1" by repeating the data. */
+	expand(newShape) {
+		if (newShape.length !== this.ndim) throw new Error(`Can't expand ${jstr(this.shape)} into ${jstr(newShape)}`);
+		for (let i = 0; i < this.ndim; i++) if (newShape[i] !== this.shape[i] && this.shape[i] !== 1) throw new Error(`Can't expand ${jstr(this.shape)} into ${jstr(newShape)}`);
+		if (this.size === 0) return View.create(newShape);
+		const mask = this.mask ? this.mask.map((m, i) => this.shape[i] === newShape[i] ? m : m[0] === 0 && m[1] === 1 ? [0, newShape[i]] : [0, 0]) : null;
+		return View.create(newShape, this.strides, this.offset, mask);
+	}
+	/** Permute the axes of an array. */
+	permute(axis) {
+		if (!isPermutation(axis, this.ndim)) throw new Error(`Invalid permutation ${jstr(axis)} of len ${this.ndim}`);
+		const newShape = axis.map((a) => this.shape[a]);
+		const newStrides = axis.map((a) => this.strides[a]);
+		const newMask = this.mask ? axis.map((a) => this.mask[a]) : null;
+		return View.create(newShape, newStrides, this.offset, newMask);
+	}
+	/** Flip (reverse) one or more axes of the view. */
+	flip(arg) {
+		if (arg.length !== this.ndim) throw new Error(`Invalid flip ${jstr(arg)} for ${jstr(this.shape)}`);
+		const strides = this.strides.slice();
+		let offset = this.offset;
+		const mask = this.mask ? this.mask.slice() : null;
+		for (let i = 0; i < this.ndim; i++) {
+			const s = this.shape[i];
+			if (arg[i]) {
+				strides[i] = -strides[i];
+				offset += (s - 1) * this.strides[i];
+				if (mask) mask[i] = [s - mask[i][1], s - mask[i][0]];
+			}
+		}
+		return View.create(this.shape, strides, offset, mask);
+	}
+	/** Reshape the view into a new shape. */
+	reshape(newShape) {
+		if (deepEqual(this.shape, newShape)) return this;
+		if (newShape.some((s) => s < 0)) throw new Error(`Reshape cannot have negative numbers ${jstr(newShape)}`);
+		if (this.size !== prod(newShape)) throw new Error(`Reshape size ${jstr(this.shape)} -> ${jstr(newShape)}`);
+		if (this.size === 0) return View.create(newShape);
+		if (newShape.length === 0 && this.mask?.some(([b, e]) => b === e)) return null;
+		if (this.contiguous) return View.create(newShape);
+		const rStrides = [];
+		const merge = mergeDims(this.shape, this.strides, this.mask);
+		let rShapeIdx = newShape.length;
+		for (let i = merge.length - 1; i >= 0; i--) {
+			let [mergedSize, newStride, realSize] = merge[i];
+			let acc = 1;
+			while (acc < mergedSize && rShapeIdx > 0) {
+				const newDim = newShape[--rShapeIdx];
+				rStrides.push(newStride * acc);
+				acc *= newDim;
+				if (acc >= realSize) newStride = 0;
+			}
+			if (acc !== mergedSize) return null;
+		}
+		const newStrides = rep(newShape.length - rStrides.length, 0).concat(rStrides.reverse());
+		if (!this.mask) return View.create(newShape, newStrides, this.offset);
+		const newMask = reshapeMask(this.mask, this.shape, newShape);
+		if (!newMask) return null;
+		let newOffset = this.offset;
+		for (let i = 0; i < this.ndim; i++) newOffset += this.strides[i] * this.mask[i][0];
+		for (let i = 0; i < newShape.length; i++) newOffset -= newStrides[i] * newMask[i][0];
+		return View.create(newShape, newStrides, newOffset, newMask);
+	}
+};
+/**
+* Find position of `offset` in each dimension within an existing shape. Like
+* `numpy.unravel_index` in behavior.
+*/
+function unravel(shape, offset) {
+	let acc = 1;
+	const idxs = [];
+	for (let i = shape.length - 1; i >= 0; i--) {
+		const d = shape[i];
+		idxs.push(Math.floor(offset / acc) % d);
+		acc *= d;
+	}
+	return idxs.reverse();
+}
+/** Generate a list of AluExp for computing unravel(). */
+function unravelAlu(shape, offset) {
+	let acc = 1;
+	const idxs = [];
+	for (let i = shape.length - 1; i >= 0; i--) {
+		const d = shape[i];
+		idxs.push(AluExp.mod(AluExp.idiv(offset, AluExp.i32(acc)), AluExp.i32(d)));
+		acc *= d;
+	}
+	return idxs.reverse();
+}
+/**
+* Array shape after applying movement operations, as a series of views.
+*
+* Each view is applied, then treated as if it were a contiguous array of its
+* shape, then used as the virtual buffer for the next view.
+*/
+var ShapeTracker = class ShapeTracker {
+	constructor(views) {
+		this.views = views;
+	}
+	/** Compose this shape tracker with another, applying it after this one. */
+	compose(other) {
+		if (this.contiguous) return other;
+		let ret = this;
+		for (const v of other.views) ret = new ShapeTracker(ret.views.concat(v)).simplify();
+		return ret;
+	}
+	static fromShape(shape) {
+		return new ShapeTracker([View.create(shape)]);
+	}
+	get contiguous() {
+		return this.views.length === 1 && this.views[0].contiguous;
+	}
+	get consecutive() {
+		return this.views.length === 1 && this.views[0].mask === null && deepEqual(this.views[0].strides, defaultStrides(this.views[0].shape));
+	}
+	get lastStrides() {
+		return this.views[this.views.length - 1].strides;
+	}
+	get shape() {
+		return this.views[this.views.length - 1].shape;
+	}
+	get size() {
+		return this.views[this.views.length - 1].size;
+	}
+	toAluExp(idxs) {
+		let [iexpr, vexpr] = this.views[this.views.length - 1].toAluExp(idxs);
+		for (let i = this.views.length - 2; i >= 0; i--) {
+			const view = this.views[i].minify();
+			const exprs = view.toAluExp(unravelAlu(view.shape, iexpr));
+			iexpr = exprs[0];
+			vexpr = AluExp.mul(vexpr, exprs[1]);
+		}
+		return [iexpr.simplify(), vexpr.simplify()];
+	}
+	simplify() {
+		const views = this.views.slice();
+		while (views.length >= 2) {
+			const newView = views[views.length - 2].compose(views[views.length - 1]);
+			if (newView === null) break;
+			views.splice(views.length - 2, 2, newView);
+		}
+		return new ShapeTracker(views);
+	}
+	pad(arg) {
+		return new ShapeTracker(applyLast(this.views, (x) => x.pad(arg)));
+	}
+	shrink(arg) {
+		return new ShapeTracker(applyLast(this.views, (x) => x.shrink(arg)));
+	}
+	expand(newShape) {
+		return new ShapeTracker(applyLast(this.views, (x) => x.expand(newShape)));
+	}
+	permute(axis) {
+		return new ShapeTracker(applyLast(this.views, (x) => x.permute(axis)));
+	}
+	flip(arg) {
+		return new ShapeTracker(applyLast(this.views, (x) => x.flip(arg)));
+	}
+	reshape(newShape) {
+		const newView = this.views[this.views.length - 1].reshape(newShape);
+		return new ShapeTracker(newView === null ? this.views.concat(View.create(newShape)) : this.views.toSpliced(this.views.length - 1, 1, newView));
+	}
+	/** Broadcast along the given new axes, then expand the shape. */
+	broadcast(newShape, axis) {
+		let st = this;
+		if (axis.length > 0) {
+			const unsqueezed = [...st.shape];
+			for (const i of axis.toSorted()) unsqueezed.splice(i, 0, 1);
+			st = st.reshape(unsqueezed);
+		}
+		return st.expand(newShape);
+	}
+};
+function applyLast(ar, f) {
+	return ar.toSpliced(ar.length - 1, 1, f(ar[ar.length - 1]));
+}
+
+//#endregion
+//#region src/tuner.ts
+/** Stores dimensions of the kernel's applied shape. Globals start at 0. */
+var TuneDims = class {
+	st;
+	outputSt;
+	groups;
+	reduce;
+	unroll;
+	upcast;
+	get end() {
+		return this.st.shape.length;
+	}
+	constructor(shape) {
+		this.st = ShapeTracker.fromShape(shape);
+		this.outputSt = ShapeTracker.fromShape(shape.slice(0, -1));
+		this.groups = this.st.shape.length - 1;
+		this.reduce = this.st.shape.length - 1;
+		this.unroll = this.st.shape.length;
+		this.upcast = this.st.shape.length;
+	}
+	applyLocal(axis, amount) {
+		if (axis >= this.groups) throw new Error("Cannot localize reduction axis");
+		const length = this.st.shape[axis];
+		if (length % amount !== 0) throw new Error(`Localize by ${amount} on axis length ${length}`);
+		if (length !== amount) {
+			this.groups++, this.reduce++, this.unroll++, this.upcast++;
+			this.st = this.st.reshape([
+				...this.st.shape.slice(0, axis),
+				length / amount,
+				amount,
+				...this.st.shape.slice(axis + 1)
+			]);
+			this.outputSt = this.outputSt.reshape([
+				...this.outputSt.shape.slice(0, axis),
+				length / amount,
+				amount,
+				...this.outputSt.shape.slice(axis + 1)
+			]);
+			axis++;
+		}
+		this.st = this.st.permute([
+			...range(axis),
+			...range(axis + 1, this.groups),
+			axis,
+			...range(this.groups, this.st.shape.length)
+		]);
+		this.outputSt = this.outputSt.permute([
+			...range(axis),
+			...range(axis + 1, this.groups),
+			axis,
+			...range(this.groups, this.outputSt.shape.length)
+		]);
+	}
+	applyUpcast(axis, amount) {
+		if (axis >= this.groups) throw new Error("Cannot upcast along reduction axis");
+		const length = this.st.shape[axis];
+		if (length % amount !== 0) throw new Error(`Upcast by ${amount} on axis length ${length}`);
+		this.st = this.st.reshape([
+			...this.st.shape.slice(0, axis),
+			length / amount,
+			amount,
+			...this.st.shape.slice(axis + 1)
+		]).permute([
+			...range(axis + 1),
+			...range(axis + 2, this.st.shape.length + 1),
+			axis + 1
+		]);
+		this.outputSt = this.outputSt.reshape([
+			...this.outputSt.shape.slice(0, axis),
+			length / amount,
+			amount,
+			...this.outputSt.shape.slice(axis + 1)
+		]).permute([
+			...range(axis + 1),
+			...range(axis + 2, this.outputSt.shape.length + 1),
+			axis + 1
+		]);
+	}
+	applyUnroll(axis, amount) {
+		if (axis < this.groups) throw new Error("Cannot unroll non-reduce axis");
+		if (axis >= this.unroll) throw new Error("Axis already unrolled");
+		const length = this.st.shape[axis];
+		if (length % amount !== 0) throw new Error(`Unroll by ${amount} on axis length ${length}`);
+		if (length === amount) {
+			this.st = this.st.permute([
+				...range(axis),
+				...range(axis + 1, this.upcast),
+				axis,
+				...range(this.upcast, this.st.shape.length)
+			]);
+			if (axis < this.reduce) this.reduce--;
+			this.unroll--;
+		} else {
+			this.st = this.st.reshape([
+				...this.st.shape.slice(0, axis),
+				length / amount,
+				amount,
+				...this.st.shape.slice(axis + 1)
+			]).permute([
+				...range(axis + 1),
+				...range(axis + 2, this.upcast + 1),
+				axis + 1,
+				...range(this.upcast + 1, this.st.shape.length + 1)
+			]);
+			this.upcast++;
+		}
+	}
+};
+/** Tuning step that does not apply any optimization. */
+function tuneNullopt(kernel) {
+	const vars = {};
+	vars.gidx = AluExp.special(DType.Int32, "gidx", kernel.size);
+	if (kernel.reduction) vars.ridx = AluExp.special(DType.Int32, "ridx", kernel.reduction.size);
+	return {
+		exp: kernel.exp.rewrite((exp) => {
+			if (exp.op === AluOp.GlobalView) {
+				const gid = exp.arg[0];
+				const st = exp.arg[1];
+				return accessorGlobal(exp.dtype, gid, st, exp.src);
+			}
+		}).substitute(vars).simplify(),
+		outputIdxExp: AluExp.special(DType.Int32, "gidx", kernel.size),
+		threadCount: kernel.size,
+		size: { reduce: kernel.reduction ? kernel.reduction.size : 0 }
+	};
+}
+/** Tuning for WebGPU kernels. */
+function tuneWebgpu(kernel) {
+	const { exp, reduction } = kernel;
+	if (!reduction) return tuneNullopt(kernel);
+	const globalIndexes = exp.collect((exp$1) => exp$1.op === AluOp.GlobalIndex);
+	if (globalIndexes.length > 0) {
+		if (DEBUG >= 4) console.info("Tuning: Found GlobalIndex ops, skipping opt.");
+		return tuneNullopt(kernel);
+	}
+	const globalViews = exp.collect((exp$1) => exp$1.op === AluOp.GlobalView);
+	if (globalViews.length === 0) {
+		if (DEBUG >= 4) console.info("Tuning: No GlobalView ops found in kernel.");
+		return tuneNullopt(kernel);
+	}
+	const shape = globalViews[0].arg[1].shape;
+	const expectedSrc = [...unravelAlu(shape.slice(0, -1), AluVar.gidx), AluVar.ridx].map((e) => e.simplify());
+	for (const gv of globalViews) if (!gv.src.length || !deepEqual(gv.src, expectedSrc)) {
+		if (DEBUG >= 4) console.info("Tuning: GlobalView src[] not consistent with reduction.");
+		return tuneNullopt(kernel);
+	}
+	if (shape[shape.length - 1] !== reduction.size) throw new Error("Invariant violation: shape doesn't match reduction size.");
+	const sts = globalViews.map((gv) => gv.arg[1]);
+	for (const st of sts) if (!deepEqual(st.shape, shape)) throw new Error("Invariant violation: GlobalView shape mismatch");
+	const dim = new TuneDims(shape);
+	const upcastedAxis = /* @__PURE__ */ new Set();
+	while (prod(dim.st.shape.slice(0, dim.groups)) >= 1024) {
+		const choices = [];
+		const composedSts = sts.map((st) => st.compose(dim.st));
+		for (let axis = 0; axis < dim.groups; axis++) for (const amount of [3, 4]) if (!upcastedAxis.has(axis) && dim.st.shape[axis] % amount === 0 && composedSts.some((st) => st.lastStrides[axis] === 0 && st.lastStrides.slice(dim.unroll).every((stride) => stride > 0))) {
+			let nonzeroStrides = 0;
+			let totalStrides = 0;
+			for (const st of composedSts) {
+				nonzeroStrides += st.lastStrides[axis] > 0 ? 1 : 0;
+				totalStrides += st.lastStrides[axis];
+			}
+			choices.push([
+				nonzeroStrides,
+				totalStrides,
+				axis,
+				amount
+			]);
+		}
+		if (choices.length > 0) {
+			choices.sort(lexCompare);
+			dim.applyUpcast(choices[0][2], choices[0][3]);
+			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)) {
+		const s = dim.st.shape[dim.unroll - 1];
+		if (s <= 32) dim.applyUnroll(dim.reduce, s);
+		else for (const splits of [4]) if (s % splits === 0) {
+			dim.applyUnroll(dim.unroll - 1, splits);
+			break;
+		}
+	}
+	for (const ax of Array.from(upcastedAxis).sort()) {
+		const s = dim.st.shape[ax];
+		for (const amount of [8, 4]) if (s % amount === 0) {
+			dim.applyLocal(ax, amount);
+			break;
+		}
+	}
+	const indices = [];
+	const addIndices = (s, exp$1) => {
+		if (s.length === 0) return;
+		else if (s.length === 1) indices.push(exp$1);
+		else indices.push(...unravelAlu(s, exp$1));
+	};
+	if (0 < dim.groups) {
+		const s = dim.st.shape.slice(0, dim.groups);
+		addIndices(s, AluExp.special(DType.Int32, "gidx", prod(s)));
+	}
+	if (dim.groups < dim.reduce) {
+		const s = dim.st.shape.slice(dim.groups, dim.reduce);
+		addIndices(s, AluExp.special(DType.Int32, "group", prod(s)));
+	}
+	if (dim.reduce <= dim.unroll) {
+		const s = dim.st.shape.slice(dim.reduce, dim.unroll);
+		addIndices(s, AluExp.special(DType.Int32, "ridx", prod(s)));
+	}
+	if (dim.unroll < dim.upcast) {
+		const s = dim.st.shape.slice(dim.unroll, dim.upcast);
+		addIndices(s, AluVar.unroll);
+	}
+	if (dim.upcast < dim.end) {
+		const s = dim.st.shape.slice(dim.upcast);
+		addIndices(s, AluVar.upcast);
+	}
+	const newExp = exp.rewrite((exp$1) => {
+		if (exp$1.op === AluOp.GlobalView) {
+			const gid = exp$1.arg[0];
+			const st = exp$1.arg[1];
+			return accessorGlobal(exp$1.dtype, gid, st.compose(dim.st), indices);
+		}
+	});
+	const outputGidx = dim.outputSt.shape.slice(0, dim.groups);
+	const outputUpcast = dim.outputSt.shape.slice(dim.groups);
+	const [outputIdxExp, _] = dim.outputSt.toAluExp([...unravelAlu(outputGidx, AluExp.special(DType.Int32, "gidx", prod(outputGidx))), ...unravelAlu(outputUpcast, AluVar.upcast)]);
+	if (prod(dim.st.shape.slice(dim.groups, dim.upcast)) !== reduction.size) throw new Error(`Invariant violation: reduction size ${reduction.size} does not match tuned dims ${JSON.stringify(dim.st.shape.slice(dim.groups, dim.upcast))}`);
+	const size = {
+		groups: prod(dim.st.shape.slice(dim.groups, dim.reduce)),
+		reduce: prod(dim.st.shape.slice(dim.reduce, dim.unroll)),
+		unroll: prod(dim.st.shape.slice(dim.unroll, dim.upcast)),
+		upcast: prod(dim.st.shape.slice(dim.upcast))
+	};
+	return {
+		exp: newExp.simplify(),
+		outputIdxExp: outputIdxExp.simplify(),
+		threadCount: kernel.size / size.upcast * size.groups,
+		size
+	};
+}
+
+//#endregion
+//#region src/backend/cpu.ts
+/** Most basic implementation of `Backend` for testing. */
+var CPUBackend = class {
+	type = "cpu";
+	maxArgs = Infinity;
+	#buffers;
+	#nextSlot;
+	constructor() {
+		this.#buffers = /* @__PURE__ */ new Map();
+		this.#nextSlot = 1;
+	}
+	malloc(size, initialData) {
+		const buffer = new ArrayBuffer(size);
+		if (initialData) {
+			if (initialData.byteLength !== size) throw new Error("initialData size does not match buffer size");
+			new Uint8Array(buffer).set(new Uint8Array(initialData));
+		}
+		const slot = this.#nextSlot++;
+		this.#buffers.set(slot, {
+			buffer,
+			ref: 1
+		});
+		return slot;
+	}
+	incRef(slot) {
+		const buffer = this.#buffers.get(slot);
+		if (!buffer) throw new SlotError(slot);
+		buffer.ref++;
+	}
+	decRef(slot) {
+		const buffer = this.#buffers.get(slot);
+		if (!buffer) throw new SlotError(slot);
+		buffer.ref--;
+		if (buffer.ref === 0) this.#buffers.delete(slot);
+	}
+	async read(slot, start, count) {
+		return this.readSync(slot, start, count);
+	}
+	readSync(slot, start, count) {
+		const buffer = this.#getBuffer(slot);
+		if (start === void 0) start = 0;
+		if (count === void 0) count = buffer.byteLength - start;
+		return buffer.slice(start, start + count);
+	}
+	async prepare(kernel) {
+		return this.prepareSync(kernel);
+	}
+	prepareSync(kernel) {
+		return new Executable(kernel, void 0);
+	}
+	dispatch({ kernel }, inputs, outputs) {
+		const { exp } = tuneNullopt(kernel);
+		const inputBuffers = inputs.map((slot) => this.#getBuffer(slot));
+		const outputBuffers = outputs.map((slot) => this.#getBuffer(slot));
+		const usedArgs = new Map(exp.collect((exp$1) => exp$1.op === AluOp.GlobalIndex).map((exp$1) => [exp$1.arg, exp$1.dtype]));
+		const inputArrays = inputBuffers.map((buf, i) => {
+			const dtype = usedArgs.get(i);
+			if (!dtype) return null;
+			return dtypedArray(dtype, buf);
+		});
+		const outputArray = dtypedArray(kernel.dtype, outputBuffers[0]);
+		const globals = (gid, bufidx) => {
+			if (gid < 0 || gid >= inputArrays.length) throw new Error("gid out of bounds: " + gid);
+			if (bufidx < 0 || bufidx >= inputArrays[gid].length) throw new Error("bufidx out of bounds: " + bufidx);
+			return inputArrays[gid][bufidx];
+		};
+		if (!kernel.reduction) for (let i = 0; i < kernel.size; i++) outputArray[i] = exp.evaluate({ gidx: i }, globals);
+		else for (let i = 0; i < kernel.size; i++) {
+			let acc = kernel.reduction.identity;
+			for (let j = 0; j < kernel.reduction.size; j++) {
+				const item = exp.evaluate({
+					gidx: i,
+					ridx: j
+				}, globals);
+				acc = kernel.reduction.evaluate(acc, item);
+			}
+			outputArray[i] = kernel.reduction.fusion.evaluate({ acc });
+		}
+	}
+	#getBuffer(slot) {
+		const buffer = this.#buffers.get(slot);
+		if (!buffer) throw new SlotError(slot);
+		return buffer.buffer;
+	}
+};
+
+//#endregion
+//#region src/backend.ts
+const devices = ["cpu", "webgpu"];
+let defaultBackend = "cpu";
+const initializedBackends = /* @__PURE__ */ new Map();
+initializedBackends.set("cpu", new CPUBackend());
+/** Set the default device backend (must be initialized). */
+function setDevice(device) {
+	if (initializedBackends.has(device)) defaultBackend = device;
+	else throw new Error(`Backend not initialized: ${device}`);
+}
+/**
+* Initialize `jax-js` library backends.
+*
+* By default, this will initialize all available backends. If one or more
+* backends is provided, only attempt to initialize those. Returns a list of
+* available backends.
+*/
+async function init(...devicesToInit) {
+	if (devicesToInit.length === 0) devicesToInit = devices;
+	const promises = [];
+	for (const device of new Set(devicesToInit)) if (!initializedBackends.has(device)) promises.push((async () => {
+		const backend = await createBackend(device);
+		if (backend) initializedBackends.set(device, backend);
+	})());
+	await Promise.all(promises);
+	return Array.from(initializedBackends.keys());
+}
+/** Create a backend, if available. Internal function called by `init()`. */
+async function createBackend(device) {
+	if (device === "cpu") return new CPUBackend();
+	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-JVpVad6g.js");
+		const importantLimits = [
+			"maxBufferSize",
+			"maxComputeInvocationsPerWorkgroup",
+			"maxComputeWorkgroupSizeX",
+			"maxComputeWorkgroupSizeY",
+			"maxComputeWorkgroupSizeZ",
+			"maxComputeWorkgroupStorageSize",
+			"maxComputeWorkgroupsPerDimension",
+			"maxStorageBufferBindingSize",
+			"maxStorageBuffersPerShaderStage",
+			"maxStorageTexturesPerShaderStage"
+		];
+		try {
+			const device$1 = await adapter.requestDevice({ requiredLimits: Object.fromEntries(importantLimits.map((feature) => [feature, adapter.limits[feature]])) });
+			return new WebGPUBackend(device$1);
+		} catch (error) {
+			console.error("Unexpected error requesting WebGPU device:", error);
+			return null;
+		}
+	} else throw new Error(`Backend not found: ${device}`);
+}
+/** Retrieve a backend that has been initialized. */
+function getBackend(device) {
+	device = device ?? defaultBackend;
+	const backend = initializedBackends.get(device);
+	if (!backend) throw new Error(`${device} backend not ready, call init() first`);
+	return backend;
+}
+var Executable = class {
+	constructor(kernel, data) {
+		this.kernel = kernel;
+		this.data = data;
+	}
+};
+var SlotError = class extends Error {
+	constructor(slot) {
+		super(`Used a buffer that is invalid or already freed: ${slot}`);
+	}
+};
+
+//#endregion
+export { AluExp, AluGroup, AluOp, AluVar, DEBUG, DType, Executable, FpHash, Kernel, PPrint, Reduction, ShapeTracker, SlotError, accessorAluExp, accessorGlobal, byteWidth, checkAxis, deepEqual, devices, dtypedArray, findPow2, getBackend, init, invertPermutation, isFloatDtype, isNumberPair, isPermutation, partitionList, prod, range, recursiveFlatten, rep, runWithCache, setDevice, strip1, toposort, tuneWebgpu, unravelAlu, unzip2, zip };