@jax-js/jax 0.0.3 → 0.0.5

package/dist/index.cjs

@@ -30,7 +30,7 @@ var __toESM = (mod, isNodeMode, target) => (target = mod != null ? __create(__ge
 }) : target, mod));
 
 //#endregion
-const require_backend = require('./backend-D2C4MJRP.cjs');
+const require_backend = require('./backend-yEU0L_ig.cjs');
 
 //#region src/tree.ts
 var tree_exports = {};
@@ -354,6 +354,8 @@ let Primitive = /* @__PURE__ */ function(Primitive$1) {
 	Primitive$1["RandomBits"] = "random_bits";
 	Primitive$1["Sin"] = "sin";
 	Primitive$1["Cos"] = "cos";
+	Primitive$1["Asin"] = "asin";
+	Primitive$1["Atan"] = "atan";
 	Primitive$1["Exp"] = "exp";
 	Primitive$1["Log"] = "log";
 	Primitive$1["Sqrt"] = "sqrt";
@@ -421,6 +423,12 @@ function sin$1(x) {
 function cos$1(x) {
 	return bind1(Primitive.Cos, [x]);
 }
+function asin$1(x) {
+	return bind1(Primitive.Asin, [x]);
+}
+function atan$1(x) {
+	return bind1(Primitive.Atan, [x]);
+}
 function exp$1(x) {
 	return bind1(Primitive.Exp, [x]);
 }
@@ -436,18 +444,16 @@ function min$1(x, y) {
 function max$1(x, y) {
 	return bind1(Primitive.Max, [x, y]);
 }
-function reduce(x, op, axis, opts) {
+function reduce(x, op, axis = null, opts) {
 	if (!require_backend.AluGroup.Reduce.has(op)) throw new TypeError(`Invalid reduce operation: ${op}`);
-	if (axis === void 0) if (x instanceof Tracer) axis = require_backend.range(x.shape.length);
-	else axis = [];
-	else if (typeof axis === "number") axis = [require_backend.checkAxis(axis, ndim$1(x))];
-	else axis = axis.map((a) => require_backend.checkAxis(a, ndim$1(x)));
+	axis = require_backend.normalizeAxis(axis, ndim$1(x));
 	const originalShape = getShape(x);
-	const result = bind1(Primitive.Reduce, [x], {
+	let result = bind1(Primitive.Reduce, [x], {
 		op,
 		axis
 	});
-	return opts?.keepDims ? broadcast(result, originalShape, axis) : result;
+	if (opts?.keepdims) result = result.reshape(originalShape.map((dim, i) => axis.includes(i) ? 1 : dim));
+	return result;
 }
 function dot$1(x, y) {
 	return bind1(Primitive.Dot, [x, y]);
@@ -493,10 +499,11 @@ function where$1(cond, x, y) {
 }
 function transpose$1(x, perm) {
 	perm = perm ? perm.map((a) => require_backend.checkAxis(a, ndim$1(x))) : require_backend.range(ndim$1(x)).reverse();
+	if (!require_backend.isPermutation(perm, ndim$1(x))) throw new Error(`Invalid transpose permutation for ${ndim$1(x)} axes: ${JSON.stringify(perm)}`);
 	return bind1(Primitive.Transpose, [x], { perm });
 }
 function broadcast(x, shape$1, axis) {
-	axis = axis.map((a) => require_backend.checkAxis(a, shape$1.length));
+	axis = require_backend.normalizeAxis(axis, shape$1.length);
 	return bind1(Primitive.Broadcast, [x], {
 		shape: shape$1,
 		axis
@@ -515,7 +522,7 @@ function reshape$1(x, shape$1) {
 	return bind1(Primitive.Reshape, [x], { shape: shape$1 });
 }
 function flip$1(x, axis) {
-	axis = axis.map((a) => require_backend.checkAxis(a, ndim$1(x)));
+	axis = require_backend.normalizeAxis(axis, ndim$1(x));
 	return bind1(Primitive.Flip, [x], { axis });
 }
 function shrink(x, slice) {
@@ -589,21 +596,49 @@ var Trace = class {
 		this.main = main;
 	}
 };
+/**
+* Broadcast shapes and promote types with casting for two avals.
+*
+* This implements the weak type behavior described in `promoteTypes()`, but not
+* implemented in that function as `weakType` is not passed.
+*/
+function promoteAvals(a, b) {
+	const shape$1 = require_backend.generalBroadcast(a.shape, b.shape);
+	const weakType = a.weakType && b.weakType;
+	let dtype;
+	if (a.weakType === b.weakType) dtype = require_backend.promoteTypes(a.dtype, b.dtype);
+	else if (a.weakType) dtype = require_backend.promoteTypes(b.dtype, require_backend.DType.Uint32);
+	else dtype = require_backend.promoteTypes(a.dtype, require_backend.DType.Uint32);
+	return new ShapedArray(shape$1, dtype, weakType);
+}
 var Tracer = class Tracer {
 	/** @ignore */
 	_trace;
 	constructor(trace) {
 		this._trace = trace;
 	}
+	/** The shape of the array. */
 	get shape() {
 		return this.aval.shape;
 	}
+	/** The total number of elements in the array. */
 	get size() {
 		return require_backend.prod(this.shape);
 	}
+	/** The dtype of elements stored in the array. */
 	get dtype() {
 		return this.aval.dtype;
 	}
+	/**
+	* Whether the array is weakly typed.
+	*
+	* Weakly typed arrays will cast to the dtype of the other operand. See
+	* `promoteTypes()` for details.
+	*/
+	get weakType() {
+		return this.aval.weakType;
+	}
+	/** The number of dimensions of the array. */
 	get ndim() {
 		return this.shape.length;
 	}
@@ -639,22 +674,20 @@ var Tracer = class Tracer {
 		return lessEqual$1(this, other);
 	}
 	/** Sum of the elements of the array over a given axis, or axes. */
-	sum(axis, opts) {
+	sum(axis = null, opts) {
 		return reduce(this, require_backend.AluOp.Add, axis, opts);
 	}
 	/** Product of the array elements over a given axis. */
-	prod(axis, opts) {
+	prod(axis = null, opts) {
 		return reduce(this, require_backend.AluOp.Mul, axis, opts);
 	}
 	/** Compute the average of the array elements along the specified axis. */
-	mean(axis, opts) {
-		if (axis === void 0) axis = require_backend.range(this.ndim);
-		else if (typeof axis === "number") axis = [require_backend.checkAxis(axis, this.ndim)];
-		else axis = axis.map((a) => require_backend.checkAxis(a, this.ndim));
-		let result = reduce(this, require_backend.AluOp.Add, axis);
-		result = result.mul(result.size / this.size);
-		if (opts?.keepDims) result = broadcast(result, this.shape, axis);
-		return result;
+	mean(axis = null, opts) {
+		axis = require_backend.normalizeAxis(axis, this.ndim);
+		const n = axis.reduce((acc, a) => acc * this.shape[a], 1);
+		if (n === 0) throw new Error("mean: cannot compute mean over zero-length axis");
+		const result = reduce(this, require_backend.AluOp.Add, axis, opts);
+		return result.mul(1 / n);
 	}
 	/** Permute the dimensions of an array. Defaults to reversing the axis order. */
 	transpose(perm) {
@@ -841,12 +874,13 @@ function getShape(x) {
 	return x instanceof Tracer ? x.shape : [];
 }
 var ShapedArray = class ShapedArray {
-	constructor(shape$1, dtype) {
+	constructor(shape$1, dtype, weakType) {
 		this.shape = shape$1;
 		this.dtype = dtype;
+		this.weakType = weakType;
 	}
 	static fromAval(aval) {
-		return new ShapedArray(aval.shape, aval.dtype);
+		return new ShapedArray(aval.shape, aval.dtype, aval.weakType);
 	}
 	get ndim() {
 		return this.shape.length;
@@ -860,7 +894,7 @@ var ShapedArray = class ShapedArray {
 };
 function getAval(x) {
 	if (x instanceof Tracer) return x.aval;
-	else if (typeof x === "boolean" || typeof x === "number") return new ShapedArray([], typeof x === "boolean" ? require_backend.DType.Bool : require_backend.DType.Float32);
+	else if (typeof x === "boolean" || typeof x === "number") return new ShapedArray([], typeof x === "boolean" ? require_backend.DType.Bool : require_backend.DType.Float32, typeof x === "boolean" ? false : true);
 	else throw new TypeError(`Unknown value: ${x}`);
 }
 function bind(prim, args, params = {}) {
@@ -1145,7 +1179,7 @@ function reshapeViews(exp$2, mapping, reduceAxis = false) {
 }
 function broadcastedJit(fn) {
 	return (nargs, exps, avals, params) => {
-		const newShape = avals.map((aval) => aval.shape).reduce(generalBroadcast);
+		const newShape = avals.map((aval) => aval.shape).reduce(require_backend.generalBroadcast);
 		exps = exps.map((exp$3) => reshapeViews(exp$3, (st) => {
 			if (!require_backend.deepEqual(st.shape, newShape)) return st.broadcast(newShape, require_backend.range(newShape.length - st.shape.length));
 		}));
@@ -1182,11 +1216,13 @@ const jitRules = {
 		const k1 = reshapeViews(keys[1], mapping);
 		const c0 = require_backend.AluExp.u32(0);
 		const c1 = require_backend.AluExp.cast(require_backend.DType.Uint32, require_backend.AluVar.gidx);
-		const exp$2 = require_backend.AluExp.threefry2x32(c0, c1, k0, k1, mode);
+		const exp$2 = require_backend.AluExp.threefry2x32(k0, k1, c0, c1, mode);
 		return new require_backend.Kernel(nargs, require_backend.prod(shape$1), exp$2);
 	},
 	[Primitive.Sin]: unopJit(require_backend.AluExp.sin),
 	[Primitive.Cos]: unopJit(require_backend.AluExp.cos),
+	[Primitive.Asin]: unopJit(require_backend.AluExp.asin),
+	[Primitive.Atan]: unopJit(require_backend.AluExp.atan),
 	[Primitive.Exp]: unopJit(require_backend.AluExp.exp),
 	[Primitive.Log]: unopJit(require_backend.AluExp.log),
 	[Primitive.Sqrt]: unopJit(require_backend.AluExp.sqrt),
@@ -1221,7 +1257,7 @@ const jitRules = {
 	[Primitive.Dot](nargs, [a, b], [as, bs]) {
 		const k1 = jitRules[Primitive.Mul](nargs, [a, b], [as, bs], {});
 		const c = k1.exp;
-		const cs = new ShapedArray(generalBroadcast(as.shape, bs.shape), c.dtype);
+		const cs = promoteAvals(as, bs);
 		return jitRules[Primitive.Reduce](nargs, [c], [cs], {
 			op: require_backend.AluOp.Add,
 			axis: [cs.ndim - 1]
@@ -1231,8 +1267,8 @@ const jitRules = {
 		const [stX, stY] = prepareConv(require_backend.ShapeTracker.fromShape(as.shape), require_backend.ShapeTracker.fromShape(bs.shape), params);
 		a = reshapeViews(a, (st) => st.compose(stX));
 		b = reshapeViews(b, (st) => st.compose(stY));
-		as = new ShapedArray(stX.shape, as.dtype);
-		bs = new ShapedArray(stY.shape, bs.dtype);
+		as = new ShapedArray(stX.shape, as.dtype, as.weakType);
+		bs = new ShapedArray(stY.shape, bs.dtype, bs.weakType);
 		return jitRules[Primitive.Dot](nargs, [a, b], [as, bs], {});
 	},
 	[Primitive.Compare]: broadcastedJit(([a, b], { op }) => aluCompare(a, b, op)),
@@ -1249,7 +1285,7 @@ const jitRules = {
 	[Primitive.Pad]: reshapeJit((st, { width }) => st.pad(width)),
 	[Primitive.Gather](nargs, [x, ...indices], [xs, ...indicesShapes], { axis, outDim }) {
 		const axisSet = new Set(axis);
-		const indexShape = indicesShapes.map((c) => c.shape).reduce(generalBroadcast);
+		const indexShape = indicesShapes.map((c) => c.shape).reduce(require_backend.generalBroadcast);
 		const finalShape = xs.shape.filter((_, i) => !axisSet.has(i));
 		finalShape.splice(outDim, 0, ...indexShape);
 		const idxAll = require_backend.unravelAlu(finalShape, require_backend.AluVar.gidx);
@@ -1285,9 +1321,10 @@ function splitGraphDataflow(backend, jaxpr) {
 		Primitive.Conv,
 		Primitive.PoolTranspose
 	];
+	const heterogeneousViewPrimitives = [Primitive.Gather, Primitive.RandomBits];
 	for (let i = jaxpr.eqns.length - 1; i >= 0; i--) {
 		const eqn = jaxpr.eqns[i];
-		if (reducePrimitives.includes(eqn.primitive) || eqn.primitive === Primitive.Gather || eqn.outBinders.some((v) => blackNodes.has(v))) {
+		if (reducePrimitives.includes(eqn.primitive) || heterogeneousViewPrimitives.includes(eqn.primitive) || eqn.outBinders.some((v) => blackNodes.has(v))) {
 			for (const v of eqn.outBinders) {
 				blackNodes.add(v);
 				p1NextBlack.set(v, v);
@@ -1417,6 +1454,7 @@ var Array$1 = class Array$1 extends Tracer {
 	static #nextId = 1001;
 	id;
 	#dtype;
+	#weakType;
 	#source;
 	#st;
 	#backend;
@@ -1428,19 +1466,22 @@ var Array$1 = class Array$1 extends Tracer {
 	* is a backend `Slot`, this constructor _takes ownership_ of the slot. It
 	* will be freed when the array is disposed.
 	*/
-	constructor(source, st, dtype, backend, pending = null) {
+	constructor(args) {
 		super(baseArrayTrace);
 		this.id = Array$1.#nextId++;
-		this.#dtype = dtype;
-		this.#source = source;
-		this.#st = st;
-		this.#backend = backend;
+		this.#dtype = args.dtype;
+		this.#weakType = args.weakType;
+		this.#source = args.source;
+		this.#st = args.st;
+		this.#backend = args.backend;
 		this.#rc = 1;
-		this.#pendingSet = new Set(pending);
+		this.#pendingSet = new Set(args.pending);
+		if (this.#pendingSet.size === 0) this.#pendingSet = null;
+		else if (this.#source instanceof require_backend.AluExp) throw new Error("internal: AluExp source cannot have pending executes");
 	}
 	/** @ignore */
 	get aval() {
-		return new ShapedArray(this.#st.shape, this.#dtype);
+		return new ShapedArray(this.#st.shape, this.#dtype, this.#weakType);
 	}
 	/** Return a simple string representation of the array's dimensions. */
 	toString() {
@@ -1452,6 +1493,17 @@ var Array$1 = class Array$1 extends Tracer {
 	#check() {
 		if (this.#rc <= 0) throw new UseAfterFreeError(this);
 	}
+	/** Construct an array, copying fields from `this`. */
+	#newArrayFrom(args) {
+		return new Array$1({
+			source: args.source ?? this.#source,
+			st: args.st ?? this.#st,
+			dtype: args.dtype ?? this.#dtype,
+			weakType: this.#weakType,
+			backend: args.backend ?? this.#backend,
+			pending: args.pending ?? this.#pending ?? void 0
+		});
+	}
 	get ref() {
 		this.#check();
 		this.#rc++;
@@ -1491,7 +1543,10 @@ var Array$1 = class Array$1 extends Tracer {
 		const pending = this.#pending;
 		for (const exe of pending) exe.updateRc(1);
 		if (typeof this.#source === "number") this.#backend.incRef(this.#source);
-		const ar = new Array$1(this.#source, st, this.#dtype, this.#backend, pending);
+		const ar = this.#newArrayFrom({
+			st,
+			pending
+		});
 		this.dispose();
 		return ar;
 	}
@@ -1540,7 +1595,11 @@ var Array$1 = class Array$1 extends Tracer {
 		pending.push(new PendingExecute(this.#backend, kernel, inputs, [output]));
 		this.dispose();
 		for (const ar of indices) ar.dispose();
-		return new Array$1(output, require_backend.ShapeTracker.fromShape(finalShape), this.#dtype, this.#backend, pending);
+		return this.#newArrayFrom({
+			source: output,
+			st: require_backend.ShapeTracker.fromShape(finalShape),
+			pending
+		});
 	}
 	/** Move axes to the rightmost dimension of the shape. */
 	#moveAxesDown(axis) {
@@ -1563,11 +1622,16 @@ var Array$1 = class Array$1 extends Tracer {
 		return this.#reshape(this.#st.permute(perm));
 	}
 	#unary(op, dtypeOutput) {
+		const weakType = !dtypeOutput && this.#weakType;
 		dtypeOutput ??= this.#dtype;
 		this.#check();
 		if (this.#source instanceof require_backend.AluExp) {
 			const exp$3 = new require_backend.AluExp(op, dtypeOutput, [this.#source]);
-			return new Array$1(exp$3.simplify(), this.#st, dtypeOutput, this.#backend);
+			return this.#newArrayFrom({
+				source: exp$3.simplify(),
+				dtype: dtypeOutput,
+				weakType
+			});
 		}
 		const indices = require_backend.unravelAlu(this.#st.shape, require_backend.AluVar.gidx);
 		const exp$2 = new require_backend.AluExp(op, dtypeOutput, [require_backend.AluExp.globalView(this.#dtype, 0, this.#st, indices)]);
@@ -1577,41 +1641,65 @@ var Array$1 = class Array$1 extends Tracer {
 		for (const exe of pending) exe.updateRc(1);
 		pending.push(new PendingExecute(this.#backend, kernel, [this.#source], [output]));
 		this.dispose();
-		return new Array$1(output, require_backend.ShapeTracker.fromShape(this.shape), dtypeOutput, this.#backend, pending);
+		return this.#newArrayFrom({
+			source: output,
+			st: require_backend.ShapeTracker.fromShape(this.shape),
+			dtype: dtypeOutput,
+			weakType,
+			pending
+		});
 	}
 	#binary(op, other) {
-		const custom = (src) => new require_backend.AluExp(op, this.#dtype, src);
+		const custom = (src) => new require_backend.AluExp(op, src[0].dtype, src);
 		return Array$1.#naryCustom(op, custom, [this, other]);
 	}
-	static #naryCustom(name, custom, arrays, { dtypeOverride, dtypeOutput, reduceAxis } = {}) {
+	static #naryCustom(name, custom, arrays, { dtypeOverride, strongTypeOutput, reduceAxis } = {}) {
 		const n = arrays.length;
 		const backend = arrays[0].#backend;
 		if (n === 0) throw new TypeError(`No inputs for ${name}`);
 		for (const ar of arrays) ar.#check();
-		let dtype;
+		let castDtype;
+		let castWeakType = true;
 		for (let i = 0; i < n; i++) {
 			if (dtypeOverride?.[i]) {
 				if (arrays[i].#dtype !== dtypeOverride[i]) throw new TypeError(`Wrong dtype in ${name}: expected ${dtypeOverride[i]}, got ${arrays[i].#dtype}`);
-			} else if (!dtype) dtype = arrays[i].#dtype;
-			else if (arrays[i].#dtype !== dtype) throw new TypeError(`Dtype mismatch in ${name}: ${dtype} vs ${arrays[i].#dtype}`);
+			} else if (castDtype === void 0) {
+				castDtype = arrays[i].#dtype;
+				castWeakType = arrays[i].#weakType;
+			} else ({dtype: castDtype, weakType: castWeakType} = promoteAvals(new ShapedArray([], castDtype, castWeakType), new ShapedArray([], arrays[i].#dtype, arrays[i].#weakType)));
 			if (arrays[i].#backend !== backend) throw new TypeError(`Backend mismatch in ${name}: ${backend.type} vs ${arrays[i].#backend.type}`);
 		}
-		dtypeOutput ??= dtype;
-		if (!dtypeOutput) throw new TypeError("nary operation with no dtype");
+		const weakType = castWeakType && !strongTypeOutput;
 		arrays = Array$1.#broadcastArrays(arrays);
 		const newShape = [...arrays[0].shape];
 		if (arrays.every((ar) => ar.#source instanceof require_backend.AluExp) && !reduceAxis) {
+			const sources = arrays.map((ar, i) => {
+				if (!dtypeOverride?.[i]) return require_backend.AluExp.cast(castDtype, ar.#source);
+				else return ar.#source;
+			});
 			if (arrays.every((ar) => require_backend.deepEqual(ar.#st, arrays[0].#st))) {
-				const exp$4 = custom(arrays.map((ar) => ar.#source));
-				return new Array$1(exp$4.simplify(), arrays[0].#st, exp$4.dtype, backend);
+				const exp$4 = custom(sources);
+				return new Array$1({
+					source: exp$4.simplify(),
+					st: arrays[0].#st,
+					dtype: exp$4.dtype,
+					weakType,
+					backend
+				});
 			}
-			const exp$3 = custom(arrays.map((ar) => {
-				const src$1 = ar.#source;
+			const exp$3 = custom(arrays.map((ar, i) => {
+				const src$1 = sources[i];
 				if (ar.#st.contiguous) return src$1;
 				return require_backend.accessorAluExp(src$1, ar.#st, require_backend.unravelAlu(newShape, require_backend.AluVar.idx));
 			}));
 			const st = require_backend.ShapeTracker.fromShape(newShape);
-			return new Array$1(exp$3.simplify(), st, exp$3.dtype, backend);
+			return new Array$1({
+				source: exp$3.simplify(),
+				st,
+				dtype: exp$3.dtype,
+				weakType,
+				backend
+			});
 		}
 		let indices;
 		if (!reduceAxis) indices = require_backend.unravelAlu(newShape, require_backend.AluVar.gidx);
@@ -1621,14 +1709,19 @@ var Array$1 = class Array$1 extends Tracer {
 		}
 		const inputs = [];
 		const src = [];
-		for (const ar of arrays) if (ar.#source instanceof require_backend.AluExp) src.push(require_backend.accessorAluExp(ar.#source, ar.#st, indices));
-		else {
-			let gid = inputs.indexOf(ar.#source);
-			if (gid === -1) {
-				gid = inputs.length;
-				inputs.push(ar.#source);
+		for (const [i, ar] of arrays.entries()) {
+			let nextSrc;
+			if (ar.#source instanceof require_backend.AluExp) nextSrc = require_backend.accessorAluExp(ar.#source, ar.#st, indices);
+			else {
+				let gid = inputs.indexOf(ar.#source);
+				if (gid === -1) {
+					gid = inputs.length;
+					inputs.push(ar.#source);
+				}
+				nextSrc = require_backend.AluExp.globalView(ar.#dtype, gid, ar.#st, indices);
 			}
-			src.push(require_backend.AluExp.globalView(ar.#dtype, gid, ar.#st, indices));
+			if (!dtypeOverride?.[i]) nextSrc = require_backend.AluExp.cast(castDtype, nextSrc);
+			src.push(nextSrc);
 		}
 		const exp$2 = custom(src);
 		let re = void 0;
@@ -1642,12 +1735,17 @@ var Array$1 = class Array$1 extends Tracer {
 		for (const exe of pending) exe.updateRc(1);
 		pending.add(new PendingExecute(backend, kernel, inputs, [output]));
 		for (const ar of arrays) ar.dispose();
-		return new Array$1(output, require_backend.ShapeTracker.fromShape(newShape), dtypeOutput, backend, pending);
+		return new Array$1({
+			source: output,
+			st: require_backend.ShapeTracker.fromShape(newShape),
+			dtype: kernel.dtype,
+			weakType,
+			backend,
+			pending
+		});
 	}
 	/** Reduce the last dimension of the array by an operation. */
 	#reduce(op) {
-		this.#check();
-		if (this.ndim === 0) throw new Error("Cannot reduce a scalar");
 		const shape$1 = this.shape;
 		const reduction = new require_backend.Reduction(this.#dtype, op, shape$1[shape$1.length - 1]);
 		const newShape = shape$1.slice(0, -1);
@@ -1666,7 +1764,11 @@ var Array$1 = class Array$1 extends Tracer {
 		for (const exe of pending) exe.updateRc(1);
 		pending.push(new PendingExecute(this.#backend, kernel, inputs, [output]));
 		this.dispose();
-		return new Array$1(output, require_backend.ShapeTracker.fromShape(newShape), this.#dtype, this.#backend, pending);
+		return this.#newArrayFrom({
+			source: output,
+			st: require_backend.ShapeTracker.fromShape(newShape),
+			pending
+		});
 	}
 	/**
 	* Normalizes this array into one backed by a `Slot`.
@@ -1702,15 +1804,15 @@ var Array$1 = class Array$1 extends Tracer {
 	}
 	#dataInline() {
 		this.#check();
-		const exp$2 = this.#source;
-		const ar = new Array$1(exp$2, this.#st, this.dtype, require_backend.getBackend("cpu"));
+		if (!(this.#source instanceof require_backend.AluExp)) throw new Error("internal: #dataInline called on non-AluExp source");
+		const ar = this.#newArrayFrom({ backend: require_backend.getBackend("cpu") });
 		this.dispose();
 		return ar.dataSync();
 	}
 	static #broadcastArrays(arrays) {
 		if (arrays.length === 0) throw new Error("Need at least one array to broadcast");
 		if (arrays.length === 1) return arrays;
-		const newShape = arrays.map((a) => a.shape).reduce(generalBroadcast);
+		const newShape = arrays.map((a) => a.shape).reduce(require_backend.generalBroadcast);
 		return arrays.map((ar) => {
 			if (require_backend.deepEqual(ar.shape, newShape)) return ar;
 			return ar.#reshape(ar.#st.broadcast(newShape, require_backend.range(newShape.length - ar.ndim)));
@@ -1739,8 +1841,11 @@ var Array$1 = class Array$1 extends Tracer {
 	*
 	* If you are mapping from `data()` or `dataSync()`, it will also trigger
 	* dispatch of operations as well.
+	*
+	* **Note:** `jax.blockUntilReady()` is a higher-level API, it calls this
+	* asynchronously for multiple arrays.
 	*/
-	async wait() {
+	async blockUntilReady() {
 		this.#check();
 		if (this.#source instanceof require_backend.AluExp) return this;
 		const pending = this.#pending;
@@ -1806,7 +1911,7 @@ var Array$1 = class Array$1 extends Tracer {
 				return [x.#binary(require_backend.AluOp.Idiv, y)];
 			},
 			[Primitive.Neg]([x]) {
-				return [zerosLike(x.ref).#binary(require_backend.AluOp.Sub, x)];
+				return [zerosLike$1(x.ref).#binary(require_backend.AluOp.Sub, x)];
 			},
 			[Primitive.Reciprocal]([x]) {
 				return [x.#unary(require_backend.AluOp.Reciprocal)];
@@ -1826,14 +1931,18 @@ var Array$1 = class Array$1 extends Tracer {
 					x.#backend.incRef(x.#source);
 					const pending = x.#pending;
 					for (const exe of pending) exe.updateRc(1);
-					const y = new Array$1(x.#source, x.#st, dtype, x.#backend, pending);
+					const y = x.#newArrayFrom({
+						dtype,
+						weakType: false,
+						pending
+					});
 					x.dispose();
 					return [y];
 				}
 			},
 			[Primitive.RandomBits]([k0, k1], { shape: shape$1, mode }) {
-				const keyShape = generalBroadcast(k0.shape, k1.shape);
-				if (!require_backend.deepEqual(generalBroadcast(keyShape, shape$1), shape$1)) throw new TypeError(`Keys of shapes ${k0.shape} and ${k1.shape} cannot be broadcast to shape ${shape$1}`);
+				const keyShape = require_backend.generalBroadcast(k0.shape, k1.shape);
+				if (!require_backend.deepEqual(require_backend.generalBroadcast(keyShape, shape$1), shape$1)) throw new TypeError(`Keys of shapes ${k0.shape} and ${k1.shape} cannot be broadcast to shape ${shape$1}`);
 				const c0 = zeros(shape$1, {
 					dtype: require_backend.DType.Uint32,
 					device: k0.device
@@ -1856,6 +1965,12 @@ var Array$1 = class Array$1 extends Tracer {
 			[Primitive.Cos]([x]) {
 				return [x.#unary(require_backend.AluOp.Cos)];
 			},
+			[Primitive.Asin]([x]) {
+				return [x.#unary(require_backend.AluOp.Asin)];
+			},
+			[Primitive.Atan]([x]) {
+				return [x.#unary(require_backend.AluOp.Atan)];
+			},
 			[Primitive.Exp]([x]) {
 				return [x.#unary(require_backend.AluOp.Exp)];
 			},
@@ -1895,7 +2010,7 @@ var Array$1 = class Array$1 extends Tracer {
 			},
 			[Primitive.Compare]([x, y], { op }) {
 				const custom = ([x$1, y$1]) => aluCompare(x$1, y$1, op);
-				return [Array$1.#naryCustom("compare", custom, [x, y], { dtypeOutput: require_backend.DType.Bool })];
+				return [Array$1.#naryCustom("compare", custom, [x, y], { strongTypeOutput: true })];
 			},
 			[Primitive.Where]([cond, x, y]) {
 				const custom = ([cond$1, x$1, y$1]) => require_backend.AluExp.where(cond$1, x$1, y$1);
@@ -1941,7 +2056,14 @@ var Array$1 = class Array$1 extends Tracer {
 				pending.splice(0, 0, ...prevPending);
 				args.forEach((x) => x.dispose());
 				return outputs.map((source, i) => {
-					return new Array$1(source, require_backend.ShapeTracker.fromShape(jaxpr.outs[i].aval.shape), jaxpr.outs[i].aval.dtype, backend, pending);
+					return new Array$1({
+						source,
+						st: require_backend.ShapeTracker.fromShape(jaxpr.outs[i].aval.shape),
+						dtype: jaxpr.outs[i].aval.dtype,
+						weakType: jaxpr.outs[i].aval.weakType,
+						backend,
+						pending
+					});
 				});
 			}
 		};
@@ -1951,33 +2073,11 @@ var Array$1 = class Array$1 extends Tracer {
 		return this.#source;
 	}
 };
-/** Construct an array from a single scalar constant. */
-function scalar(value, { dtype, device } = {}) {
-	if (typeof value === "number") {
-		dtype ??= require_backend.DType.Float32;
-		if (![
-			require_backend.DType.Float32,
-			require_backend.DType.Float16,
-			require_backend.DType.Int32,
-			require_backend.DType.Uint32
-		].includes(dtype)) throw new TypeError(`Mismatched dtype for scalar ${value}`);
-	} else if (typeof value === "boolean") {
-		dtype ??= require_backend.DType.Bool;
-		if (![
-			require_backend.DType.Float32,
-			require_backend.DType.Float16,
-			require_backend.DType.Int32,
-			require_backend.DType.Uint32,
-			require_backend.DType.Bool
-		].includes(dtype)) throw new TypeError(`Mismatched dtype for scalar ${value}`);
-	} else throw new TypeError(`Invalid type for scalar ${value}`);
-	return new Array$1(require_backend.AluExp.const(dtype, value), require_backend.ShapeTracker.fromShape([]), dtype, require_backend.getBackend(device));
-}
 /** Constructor for creating a new array from data. */
 function array(values, { shape: shape$1, dtype, device } = {}) {
 	if (values instanceof Tracer) {
 		if (shape$1 && !require_backend.deepEqual(values.shape, shape$1)) values = values.reshape(shape$1);
-		if (dtype && values.dtype !== dtype) throw new Error("array astype not implemented yet");
+		if (dtype && values.dtype !== dtype) values = values.astype(dtype);
 		return values;
 	} else if (ArrayBuffer.isView(values)) return arrayFromData(values, shape$1 ?? [values.length], {
 		dtype,
@@ -1999,6 +2099,10 @@ function array(values, { shape: shape$1, dtype, device } = {}) {
 			dtype,
 			device
 		});
+		if (size$1 === 1) return full(shape$1, flat[0], {
+			dtype,
+			device
+		});
 		if (typeof flat[0] === "boolean") {
 			dtype = dtype ?? require_backend.DType.Bool;
 			const data = new Int32Array(flat.map((x) => x ? 1 : 0));
@@ -2007,46 +2111,51 @@ function array(values, { shape: shape$1, dtype, device } = {}) {
 				device
 			});
 		} else {
+			const weakType = dtype == void 0;
 			dtype = dtype ?? require_backend.DType.Float32;
 			const data = require_backend.dtypedJsArray(dtype, flat);
 			return arrayFromData(data, shape$1, {
 				dtype,
 				device
-			});
+			}, weakType);
 		}
 	}
 }
-function arrayFromData(data, shape$1, { dtype, device } = {}) {
+function arrayFromData(data, shape$1, { dtype, device }, weakType = false) {
+	if (data instanceof Float32Array) {
+		if (dtype && dtype !== require_backend.DType.Float32) throw new Error("Float32Array must have float32 type");
+		dtype ??= require_backend.DType.Float32;
+	} else if (data instanceof Int32Array) {
+		if (dtype && dtype !== require_backend.DType.Int32 && dtype !== require_backend.DType.Bool) throw new Error("Int32Array must have int32 or bool type");
+		dtype ??= require_backend.DType.Int32;
+	} else if (data instanceof Uint32Array) {
+		if (dtype && dtype !== require_backend.DType.Uint32) throw new Error("Uint32Array must have uint32 type");
+		dtype ??= require_backend.DType.Uint32;
+	} else if (data instanceof Float16Array) {
+		if (dtype && dtype !== require_backend.DType.Float16) throw new Error("Float16Array must have float16 type");
+		dtype ??= require_backend.DType.Float16;
+	} else throw new Error("Unsupported data array type: " + data.constructor.name);
 	if (data.length < inlineArrayLimit) {
 		let allEqual = true;
 		for (let i = 1; i < data.length; i++) if (data[i] !== data[0]) {
 			allEqual = false;
 			break;
 		}
-		if (allEqual) return full(shape$1, data[0], {
-			dtype,
-			device
-		});
+		if (allEqual) {
+			const sa = new ShapedArray(shape$1, dtype, weakType);
+			return fullInternal(sa, data[0], device);
+		}
 	}
 	const backend = require_backend.getBackend(device);
-	if (ArrayBuffer.isView(data)) {
-		const buf = new Uint8Array(data.buffer, data.byteOffset, data.byteLength);
-		if (data instanceof Float32Array) {
-			if (dtype && dtype !== require_backend.DType.Float32) throw new Error("Float32Array must have float32 type");
-			dtype ??= require_backend.DType.Float32;
-		} else if (data instanceof Int32Array) {
-			if (dtype && dtype !== require_backend.DType.Int32 && dtype !== require_backend.DType.Bool) throw new Error("Int32Array must have int32 or bool type");
-			dtype ??= require_backend.DType.Int32;
-		} else if (data instanceof Uint32Array) {
-			if (dtype && dtype !== require_backend.DType.Uint32) throw new Error("Uint32Array must have uint32 type");
-			dtype ??= require_backend.DType.Uint32;
-		} else if (data instanceof Float16Array) {
-			if (dtype && dtype !== require_backend.DType.Float16) throw new Error("Float16Array must have float16 type");
-			dtype ??= require_backend.DType.Float16;
-		} else throw new Error("Unsupported data array type: " + data.constructor.name);
-		const slot = backend.malloc(data.byteLength, buf);
-		return new Array$1(slot, require_backend.ShapeTracker.fromShape(shape$1), dtype, backend);
-	} else throw new Error("Unsupported data type: " + data.constructor.name);
+	const buf = new Uint8Array(data.buffer, data.byteOffset, data.byteLength);
+	const slot = backend.malloc(data.byteLength, buf);
+	return new Array$1({
+		source: slot,
+		st: require_backend.ShapeTracker.fromShape(shape$1),
+		dtype,
+		weakType,
+		backend
+	});
 }
 function dataToJs(dtype, data, shape$1) {
 	if (shape$1.length === 0) return dtype === require_backend.DType.Bool ? Boolean(data[0]) : data[0];
@@ -2062,7 +2171,7 @@ function dataToJs(dtype, data, shape$1) {
 /** If x is a value, lift it into an array, otherwise leave it be. */
 function pureArray(x) {
 	if (x instanceof Tracer) return x;
-	else return scalar(x);
+	else return array(x);
 }
 var EvalTrace = class extends Trace {
 	pure = (x) => pureArray(x);
@@ -2073,20 +2182,27 @@ var EvalTrace = class extends Trace {
 };
 const baseArrayTrace = new EvalTrace(newMain(EvalTrace, null));
 const implRules = Array$1._implRules();
-function zerosLike(val, dtype) {
-	const aval = getAval(val);
-	if (val instanceof Tracer) val.dispose();
-	return zeros(aval.shape, { dtype: dtype ?? aval.dtype });
+function fullInternal(aval, fillValue, device) {
+	return new Array$1({
+		source: require_backend.AluExp.const(aval.dtype, fillValue),
+		st: require_backend.ShapeTracker.fromShape(aval.shape),
+		dtype: aval.dtype,
+		weakType: aval.weakType,
+		backend: require_backend.getBackend(device)
+	});
 }
-function onesLike(val, dtype) {
-	const aval = getAval(val);
-	if (val instanceof Tracer) val.dispose();
-	return ones(aval.shape, { dtype: dtype ?? aval.dtype });
+function zerosLike$1(val, dtype) {
+	return fullLike(val, 0, dtype);
+}
+function onesLike$1(val, dtype) {
+	return fullLike(val, 1, dtype);
 }
 function fullLike(val, fillValue, dtype) {
 	const aval = getAval(val);
 	if (val instanceof Tracer) val.dispose();
-	return full(aval.shape, fillValue, { dtype: dtype ?? aval.dtype });
+	if (fillValue instanceof Tracer) throw new Error("numpy.fullLike() with array argument not implemented yet");
+	const sa = new ShapedArray(aval.shape, dtype ?? aval.dtype, aval.weakType);
+	return fullInternal(sa, fillValue);
 }
 /** Return a new array of given shape and type, filled with zeros. */
 function zeros(shape$1, { dtype, device } = {}) {
@@ -2104,19 +2220,14 @@ function ones(shape$1, { dtype, device } = {}) {
 }
 /** Return a new array of given shape and type, filled with `fill_value`. */
 function full(shape$1, fillValue, { dtype, device } = {}) {
-	let source;
-	if (typeof fillValue === "number") {
-		dtype = dtype ?? require_backend.DType.Float32;
-		source = require_backend.AluExp.const(dtype, fillValue);
-	} else if (typeof fillValue === "bigint") {
-		dtype = dtype ?? require_backend.DType.Int32;
-		source = require_backend.AluExp.const(dtype, Number(fillValue));
-	} else if (typeof fillValue === "boolean") {
+	let weakType = dtype == void 0;
+	if (typeof fillValue === "number") dtype = dtype ?? require_backend.DType.Float32;
+	else if (typeof fillValue === "boolean") {
 		dtype = dtype ?? require_backend.DType.Bool;
-		source = require_backend.AluExp.const(dtype, fillValue ? 1 : 0);
+		weakType = false;
 	} else if (fillValue instanceof Tracer) throw new Error("numpy.full() with array argument not implemented yet");
 	else throw new TypeError(`Invalid type for full: ${fillValue}`);
-	return new Array$1(source, require_backend.ShapeTracker.fromShape(shape$1), dtype ?? require_backend.DType.Float32, require_backend.getBackend(device));
+	return fullInternal(new ShapedArray(shape$1, dtype, weakType), fillValue, device);
 }
 /**
 * Create an identity matrix.
@@ -2126,6 +2237,7 @@ function full(shape$1, fillValue, { dtype, device } = {}) {
 */
 function eye(numRows, numCols, { dtype, device } = {}) {
 	numCols = numCols ?? numRows;
+	const weakType = dtype == void 0;
 	dtype = dtype ?? require_backend.DType.Float32;
 	if (numCols < numRows) {
 		const arr = eye(numCols, numRows, {
@@ -2139,9 +2251,15 @@ function eye(numRows, numCols, { dtype, device } = {}) {
 		device
 	});
 	const exp$2 = require_backend.AluExp.cmplt(require_backend.AluExp.mod(require_backend.AluVar.idx, require_backend.AluExp.i32(numCols + 1)), require_backend.AluExp.i32(1));
-	return new Array$1(require_backend.AluExp.cast(dtype, exp$2), require_backend.ShapeTracker.fromShape([numRows, numCols]), dtype, require_backend.getBackend(device));
+	return new Array$1({
+		source: require_backend.AluExp.cast(dtype, exp$2),
+		st: require_backend.ShapeTracker.fromShape([numRows, numCols]),
+		dtype,
+		weakType,
+		backend: require_backend.getBackend(device)
+	});
 }
-/** Return the identity array, with ones on the main diagonal. */
+/** Return the identity matrix, with ones on the main diagonal. */
 function identity$1(n, { dtype, device } = {}) {
 	return eye(n, n, {
 		dtype,
@@ -2176,7 +2294,13 @@ function arange(start, stop, step = 1, { dtype, device } = {}) {
 	});
 	const exp$2 = require_backend.AluExp.add(require_backend.AluExp.const(dtype, start), require_backend.AluExp.mul(require_backend.AluExp.cast(dtype, require_backend.AluVar.idx), require_backend.AluExp.const(dtype, step)));
 	const st = require_backend.ShapeTracker.fromShape([size$1]);
-	return new Array$1(exp$2, st, dtype, require_backend.getBackend(device));
+	return new Array$1({
+		source: exp$2,
+		st,
+		dtype,
+		weakType: false,
+		backend: require_backend.getBackend(device)
+	});
 }
 /**
 * Return evenly spaced numbers over a specified interval.
@@ -2194,10 +2318,10 @@ function linspace(start, stop, num = 50, endpoint = true, { dtype, device } = {}
 		dtype,
 		device
 	});
-	else if (num === 1) return scalar(start, {
+	else if (num === 1) return full([1], start, {
 		dtype,
 		device
-	}).reshape([1]);
+	});
 	else if (start === stop) return full([num], start, {
 		dtype,
 		device
@@ -2206,7 +2330,13 @@ function linspace(start, stop, num = 50, endpoint = true, { dtype, device } = {}
 	const denom = endpoint ? num - 1 : num;
 	const exp$2 = require_backend.AluExp.cast(dtype, require_backend.AluExp.add(require_backend.AluExp.f32(start), require_backend.AluExp.mul(require_backend.AluExp.f32(delta / denom), require_backend.AluExp.cast(require_backend.DType.Float32, require_backend.AluVar.idx))));
 	const st = require_backend.ShapeTracker.fromShape([num]);
-	return new Array$1(exp$2, st, dtype, require_backend.getBackend(device));
+	return new Array$1({
+		source: exp$2,
+		st,
+		dtype,
+		weakType: false,
+		backend: require_backend.getBackend(device)
+	});
 }
 function aluCompare(a, b, op) {
 	switch (op) {
@@ -2218,35 +2348,6 @@ function aluCompare(a, b, op) {
 		case CompareOp.LessEqual: return require_backend.AluExp.add(require_backend.AluExp.cmplt(a, b), require_backend.AluExp.cmpne(a, b).not());
 	}
 }
-/**
-* 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();
-}
 
 //#endregion
 //#region node_modules/.pnpm/@oxc-project+runtime@0.78.0/node_modules/@oxc-project/runtime/src/helpers/usingCtx.js
@@ -2326,13 +2427,15 @@ var Var = class Var {
 };
 /** Literal in a Jaxpr expression. Currently, only scalars are supported. */
 var Lit = class {
-	dtype;
 	value;
 	aval;
-	constructor(dtype, value) {
-		this.dtype = dtype;
+	get dtype() {
+		return this.aval.dtype;
+	}
+	constructor(aval, value) {
+		if (aval.shape.length !== 0) throw new Error(`internal: Lit must be a scalar`);
 		this.value = value;
-		this.aval = new ShapedArray([], dtype);
+		this.aval = ShapedArray.fromAval(aval);
 	}
 };
 function atomIsLit(atom, literal) {
@@ -2421,16 +2524,19 @@ var Jaxpr = class Jaxpr {
 			varIds.set(v, require_backend.FpHash.hash(id, v.aval.dtype, ...v.aval.shape));
 			return id;
 		};
-		hasher.update(this.inBinders.length, ...this.inBinders.map(vi));
-		hasher.update(this.eqns.length, ...this.eqns.flatMap((eqn) => [
-			eqn.primitive,
-			eqn.inputs.length,
-			...eqn.inputs.map((x) => x instanceof Var ? vi(x) : x.value),
-			JSON.stringify(eqn.params),
-			eqn.outBinders.length,
-			...eqn.outBinders.map(vi)
-		]));
-		hasher.update(this.outs.length, ...this.outs.map((x) => x instanceof Var ? vi(x) : x.value));
+		hasher.update(this.inBinders.length);
+		for (const x of this.inBinders) hasher.update(vi(x));
+		hasher.update(this.eqns.length);
+		for (const eqn of this.eqns) {
+			hasher.update(eqn.primitive);
+			hasher.update(eqn.inputs.length);
+			for (const x of eqn.inputs) hasher.update(x instanceof Var ? vi(x) : x.value);
+			hasher.update(JSON.stringify(eqn.params));
+			hasher.update(eqn.outBinders.length);
+			for (const x of eqn.outBinders) hasher.update(vi(x));
+		}
+		hasher.update(this.outs.length);
+		for (const x of this.outs) hasher.update(x instanceof Var ? vi(x) : x.value);
 		return this.#hash = hasher.value;
 	}
 	hash(state) {
@@ -2453,21 +2559,26 @@ var Jaxpr = class Jaxpr {
 				const c = eqn.outBinders[0];
 				if (atomIsLit(a, 0)) context.set(c, b);
 				else if (atomIsLit(b, 0)) context.set(c, a);
-				else if (atomIsLit(a) && atomIsLit(b)) context.set(c, new Lit(a.dtype, a.dtype === require_backend.DType.Bool ? Math.min(a.value + b.value, 1) : a.value + b.value));
+				else if (atomIsLit(a) && atomIsLit(b)) context.set(c, new Lit(promoteAvals(a.aval, b.aval), a.dtype === require_backend.DType.Bool ? Math.min(a.value + b.value, 1) : a.value + b.value));
+				else newEqns.push(eqn);
+			} else if (eqn.primitive === Primitive.Neg) {
+				const [a] = inputs;
+				const c = eqn.outBinders[0];
+				if (atomIsLit(a)) context.set(c, new Lit(a.aval, -a.value));
 				else newEqns.push(eqn);
 			} else if (eqn.primitive === Primitive.Mul) {
 				const [a, b] = inputs;
 				const c = eqn.outBinders[0];
 				if (atomIsLit(a, 1)) context.set(c, b);
 				else if (atomIsLit(b, 1)) context.set(c, a);
-				else if (atomIsLit(a) && atomIsLit(b)) context.set(c, new Lit(a.dtype, a.value * b.value));
+				else if (atomIsLit(a) && atomIsLit(b)) context.set(c, new Lit(promoteAvals(a.aval, b.aval), a.value * b.value));
 				else newEqns.push(eqn);
 			} else if (eqn.primitive === Primitive.Idiv) {
 				const [a, b] = inputs;
 				const c = eqn.outBinders[0];
 				if (atomIsLit(b, 1)) context.set(c, a);
 				else newEqns.push(eqn);
-			} else if ((eqn.primitive === Primitive.Broadcast || eqn.primitive === Primitive.Reshape) && require_backend.deepEqual(eqn.params.shape, eqn.inputs[0].aval.shape)) context.set(eqn.outBinders[0], eqn.inputs[0]);
+			} else if ((eqn.primitive === Primitive.Broadcast || eqn.primitive === Primitive.Reshape) && require_backend.deepEqual(eqn.params.shape, eqn.inputs[0].aval.shape) || eqn.primitive === Primitive.Transpose && eqn.params.perm.every((p, i) => p === i) || eqn.primitive === Primitive.Flip && eqn.params.axis.length === 0 || eqn.primitive === Primitive.Shrink && eqn.params.slice.every(([s, e$2], i) => s === 0 && e$2 === eqn.inputs[0].aval.shape[i]) || eqn.primitive === Primitive.Pad && eqn.params.width.every(([w0, w1]) => w0 === 0 && w1 === 0)) context.set(eqn.outBinders[0], eqn.inputs[0]);
 			else newEqns.push(eqn);
 		}
 		const outs = this.outs.map((x) => x instanceof Var ? context.get(x) ?? x : x);
@@ -2558,7 +2669,7 @@ function evalJaxpr(jaxpr, args) {
 		if (x instanceof Var) {
 			remainingRefs.set(x, (remainingRefs.get(x) ?? 0) - 1);
 			return env.get(x);
-		} else return scalar(x.value, { dtype: x.dtype });
+		} else return array(x.value, { dtype: x.dtype });
 	};
 	const write = (v, val) => {
 		if (env.has(v)) throw new Error(`Variable already bound: ${v}`);
@@ -2617,7 +2728,7 @@ var JaxprTrace = class extends Trace {
 		let tracer = this.builder.constTracers.get(val);
 		if (tracer === void 0) {
 			tracer = this.builder.newTracer(this, ShapedArray.fromAval(getAval(val)));
-			this.builder.addConst(tracer, val instanceof Tracer ? val.ref : scalar(val));
+			this.builder.addConst(tracer, val instanceof Tracer ? val.ref : array(val));
 		}
 		return tracer;
 	}
@@ -2686,7 +2797,7 @@ function _inlineLiterals(jaxpr, consts) {
 	const newConsts = [];
 	for (let i = 0; i < consts.length; i++) if (ndim$1(consts[i]) === 0 && consts[i] instanceof Array$1) {
 		const ar = consts[i];
-		literals.set(jaxpr.inBinders[i], new Lit(ar.dtype, ar.dataSync()[0]));
+		literals.set(jaxpr.inBinders[i], new Lit(ar.aval, ar.dataSync()[0]));
 	} else {
 		constBinders.push(jaxpr.inBinders[i]);
 		newConsts.push(consts[i]);
@@ -2699,13 +2810,12 @@ function _inlineLiterals(jaxpr, consts) {
 }
 function binopAbstractEval([x, y]) {
 	if (!(x instanceof ShapedArray) || !(y instanceof ShapedArray)) throw new TypeError("binopAbstractEval expects ShapedArray inputs");
-	if (x.dtype !== y.dtype) throw new TypeError(`Mismatched dtypes: ${x.dtype} vs ${y.dtype}`);
-	return [new ShapedArray(generalBroadcast(x.shape, y.shape), x.dtype)];
+	return [promoteAvals(x, y)];
 }
 function compareAbstractEval([x, y]) {
 	if (!(x instanceof ShapedArray) || !(y instanceof ShapedArray)) throw new TypeError("compareAbstractEval expects ShapedArray inputs");
-	if (x.dtype !== y.dtype) throw new TypeError(`Mismatched dtypes: ${x.dtype} vs ${y.dtype}`);
-	return [new ShapedArray(generalBroadcast(x.shape, y.shape), require_backend.DType.Bool)];
+	const aval = promoteAvals(x, y);
+	return [new ShapedArray(aval.shape, require_backend.DType.Bool, false)];
 }
 function vectorizedUnopAbstractEval([x]) {
 	return [ShapedArray.fromAval(x)];
@@ -2718,21 +2828,23 @@ const abstractEvalRules = {
 	[Primitive.Reciprocal]: vectorizedUnopAbstractEval,
 	[Primitive.StopGradient]: vectorizedUnopAbstractEval,
 	[Primitive.Cast]([x], { dtype }) {
-		return [new ShapedArray(x.shape, dtype)];
+		return [new ShapedArray(x.shape, dtype, false)];
 	},
 	[Primitive.Bitcast]([x], { dtype }) {
 		if (x.dtype === require_backend.DType.Bool || dtype === require_backend.DType.Bool) throw new TypeError("Bitcast to/from bool is not allowed");
 		if (require_backend.byteWidth(x.dtype) !== require_backend.byteWidth(dtype)) throw new TypeError(`Bitcast from ${x.dtype} to ${dtype} with different byte width`);
-		return [new ShapedArray(x.shape, dtype)];
+		return [new ShapedArray(x.shape, dtype, false)];
 	},
 	[Primitive.RandomBits]([k0, k1], { shape: shape$1 }) {
 		if (k0.dtype !== require_backend.DType.Uint32 || k1.dtype !== require_backend.DType.Uint32) throw new TypeError(`RandomBits requires uint32 keys, got ${k0.dtype} and ${k1.dtype}`);
-		const keyShape = generalBroadcast(k0.shape, k1.shape);
-		if (!require_backend.deepEqual(generalBroadcast(keyShape, shape$1), shape$1)) throw new TypeError(`Keys of shapes ${k0.shape} and ${k1.shape} cannot be broadcast to shape ${shape$1}`);
-		return [new ShapedArray(shape$1, require_backend.DType.Uint32)];
+		const keyShape = require_backend.generalBroadcast(k0.shape, k1.shape);
+		if (!require_backend.deepEqual(require_backend.generalBroadcast(keyShape, shape$1), shape$1)) throw new TypeError(`Keys of shapes ${k0.shape} and ${k1.shape} cannot be broadcast to shape ${shape$1}`);
+		return [new ShapedArray(shape$1, require_backend.DType.Uint32, false)];
 	},
 	[Primitive.Sin]: vectorizedUnopAbstractEval,
 	[Primitive.Cos]: vectorizedUnopAbstractEval,
+	[Primitive.Asin]: vectorizedUnopAbstractEval,
+	[Primitive.Atan]: vectorizedUnopAbstractEval,
 	[Primitive.Exp]: vectorizedUnopAbstractEval,
 	[Primitive.Log]: vectorizedUnopAbstractEval,
 	[Primitive.Sqrt]: vectorizedUnopAbstractEval,
@@ -2741,55 +2853,54 @@ const abstractEvalRules = {
 	[Primitive.Reduce]([x], { axis }) {
 		const axisSet = new Set(axis);
 		const newShape = x.shape.filter((_, i) => !axisSet.has(i));
-		return [new ShapedArray(newShape, x.dtype)];
+		return [new ShapedArray(newShape, x.dtype, x.weakType)];
 	},
 	[Primitive.Pool]([x], { window, strides }) {
 		const shape$1 = checkPoolShape(x.shape, window, strides);
-		return [new ShapedArray(shape$1, x.dtype)];
+		return [new ShapedArray(shape$1, x.dtype, x.weakType)];
 	},
 	[Primitive.PoolTranspose]([x], { inShape, window, strides }) {
 		const shape$1 = checkPoolShape(inShape, window, strides);
 		if (!require_backend.deepEqual(shape$1, x.shape)) throw new TypeError(`PoolTranspose shape mismatch: expected ${JSON.stringify(shape$1)}, got ${JSON.stringify(x.shape)}`);
-		return [new ShapedArray(inShape, x.dtype)];
+		return [new ShapedArray(inShape, x.dtype, x.weakType)];
 	},
 	[Primitive.Dot]([x, y]) {
-		if (x.dtype !== y.dtype) throw new TypeError(`Dot dtype mismatch, got ${x.dtype} vs ${y.dtype}`);
 		if (x.ndim === 0 && y.ndim === 0) throw new TypeError("Dot requires at least 1D inputs");
-		const shape$1 = generalBroadcast(x.shape, y.shape);
+		const { shape: shape$1, dtype, weakType } = promoteAvals(x, y);
 		shape$1.splice(-1, 1);
-		return [new ShapedArray(shape$1, x.dtype)];
+		return [new ShapedArray(shape$1, dtype, weakType)];
 	},
 	[Primitive.Conv]([lhs, rhs], params) {
-		if (lhs.dtype !== rhs.dtype) throw new TypeError(`Conv dtype mismatch, got ${lhs.dtype} vs ${rhs.dtype}`);
+		const { dtype, weakType } = promoteAvals(new ShapedArray([], lhs.dtype, lhs.weakType), new ShapedArray([], rhs.dtype, rhs.weakType));
 		const shape$1 = checkConvShape(lhs.shape, rhs.shape, params);
-		return [new ShapedArray(shape$1, lhs.dtype)];
+		return [new ShapedArray(shape$1, dtype, weakType)];
 	},
 	[Primitive.Compare]: compareAbstractEval,
 	[Primitive.Where]([cond, x, y]) {
 		if (cond.dtype !== require_backend.DType.Bool) throw new TypeError(`Condition must be boolean, got ${cond.dtype}`);
-		if (x.dtype !== y.dtype) throw new TypeError(`Mismatched dtypes: ${x.dtype} vs ${y.dtype}`);
-		const shape$1 = generalBroadcast(cond.shape, generalBroadcast(x.shape, y.shape));
-		return [new ShapedArray(shape$1, x.dtype)];
+		const xy = promoteAvals(x, y);
+		const shape$1 = require_backend.generalBroadcast(cond.shape, xy.shape);
+		return [new ShapedArray(shape$1, xy.dtype, xy.weakType)];
 	},
 	[Primitive.Transpose]([x], { perm }) {
-		return [new ShapedArray(perm.map((i) => x.shape[i]), x.dtype)];
+		return [new ShapedArray(perm.map((i) => x.shape[i]), x.dtype, x.weakType)];
 	},
 	[Primitive.Broadcast]([x], { shape: shape$1 }) {
-		return [new ShapedArray(shape$1, x.dtype)];
+		return [new ShapedArray(shape$1, x.dtype, x.weakType)];
 	},
 	[Primitive.Reshape]([x], { shape: shape$1 }) {
-		return [new ShapedArray(shape$1, x.dtype)];
+		return [new ShapedArray(shape$1, x.dtype, x.weakType)];
 	},
 	[Primitive.Flip]([x], _) {
-		return [new ShapedArray(x.shape, x.dtype)];
+		return [ShapedArray.fromAval(x)];
 	},
 	[Primitive.Shrink]([x], { slice }) {
 		const newShape = slice.map((s) => s[1] - s[0]);
-		return [new ShapedArray(newShape, x.dtype)];
+		return [new ShapedArray(newShape, x.dtype, x.weakType)];
 	},
 	[Primitive.Pad]([x], { width }) {
 		const newShape = x.shape.map((dim, i) => dim + width[i][0] + width[i][1]);
-		return [new ShapedArray(newShape, x.dtype)];
+		return [new ShapedArray(newShape, x.dtype, x.weakType)];
 	},
 	[Primitive.Gather]([x, ...indices], { axis, outDim }) {
 		for (const a of indices) if (a.dtype !== require_backend.DType.Int32 && a.dtype !== require_backend.DType.Uint32) throw new TypeError(`Gather indices must be Int32 or Uint32, got ${a.dtype}`);
@@ -2799,10 +2910,10 @@ const abstractEvalRules = {
 		if (outDim < 0 || outDim > x.shape.length - axis.length) throw new TypeError("Gather outDim out of bounds");
 		const axisSet = new Set(axis);
 		if (axisSet.size !== axis.length) throw new TypeError("Gather axes are not unique");
-		const gatherShape = indices.reduce((shape$1, a) => generalBroadcast(shape$1, a.shape), []);
+		const gatherShape = indices.reduce((shape$1, a) => require_backend.generalBroadcast(shape$1, a.shape), []);
 		const newShape = x.shape.filter((_, i) => !axisSet.has(i));
 		newShape.splice(outDim, 0, ...gatherShape);
-		return [new ShapedArray(newShape, x.dtype)];
+		return [new ShapedArray(newShape, x.dtype, x.weakType)];
 	},
 	[Primitive.JitCall](args, { jaxpr }) {
 		const { inTypes, outTypes } = typecheckJaxpr(jaxpr);
@@ -2860,7 +2971,7 @@ function makeJaxpr$1(f, opts) {
 function jit$1(f, opts) {
 	const cache = /* @__PURE__ */ new Map();
 	const staticArgnums = new Set(opts?.staticArgnums ?? []);
-	return ((...args) => {
+	const result = ((...args) => {
 		const [staticArgs, dynamicArgs] = splitIdx(args, staticArgnums);
 		const [argsFlat, inTree] = flatten(dynamicArgs);
 		const avalsInFlat = argsFlat.map((x) => ShapedArray.fromAval(getAval(x)));
@@ -2869,11 +2980,16 @@ function jit$1(f, opts) {
 		const cacheKey = JSON.stringify(jaxprArgs);
 		const { jaxpr, consts, treedef: outTree } = require_backend.runWithCache(cache, cacheKey, () => makeJaxpr$1(f, opts)(...jaxprArgs));
 		const outs = bind(Primitive.JitCall, [...consts.map((c) => c.ref), ...argsFlat], {
+			name: f.name || "closure",
 			jaxpr,
 			numConsts: consts.length
 		});
 		return unflatten(outTree, outs);
 	});
+	result.dispose = () => {
+		for (const { consts } of cache.values()) for (const c of consts) c.dispose();
+	};
+	return result;
 }
 
 //#endregion
@@ -2905,7 +3021,7 @@ var JVPTrace = class extends Trace {
 		return this.lift(pureArray(val));
 	}
 	lift(val) {
-		return new JVPTracer(this, val, zerosLike(val.ref));
+		return new JVPTracer(this, val, zerosLike$1(val.ref));
 	}
 	processPrimitive(primitive, tracers, params) {
 		const [primalsIn, tangentsIn] = require_backend.unzip2(tracers.map((x) => [x.primal, x.tangent]));
@@ -2936,7 +3052,7 @@ function zeroTangentsJvp(primitive) {
 	return (primals, tangents, params) => {
 		for (const t of tangents) t.dispose();
 		const ys = bind(primitive, primals, params);
-		return [ys, ys.map((y) => zerosLike(y.ref))];
+		return [ys, ys.map((y) => zerosLike$1(y.ref))];
 	};
 }
 const jvpRules = {
@@ -2954,13 +3070,13 @@ const jvpRules = {
 		if (require_backend.isFloatDtype(dtype) && require_backend.isFloatDtype(x.dtype)) return [[cast(x, dtype)], [cast(dx, dtype)]];
 		else {
 			dx.dispose();
-			return [[cast(x.ref, dtype)], [zerosLike(x)]];
+			return [[cast(x.ref, dtype)], [zerosLike$1(x)]];
 		}
 	},
 	[Primitive.Bitcast]([x], [dx], { dtype }) {
 		if (x.dtype === dtype) return [[x], [dx]];
 		dx.dispose();
-		return [[bitcast(x.ref, dtype)], [zerosLike(x)]];
+		return [[bitcast(x.ref, dtype)], [zerosLike$1(x)]];
 	},
 	[Primitive.RandomBits]: zeroTangentsJvp(Primitive.RandomBits),
 	[Primitive.Sin]([x], [dx]) {
@@ -2969,6 +3085,14 @@ const jvpRules = {
 	[Primitive.Cos]([x], [dx]) {
 		return [[cos$1(x.ref)], [neg(sin$1(x)).mul(dx)]];
 	},
+	[Primitive.Asin]([x], [dx]) {
+		const denom = sqrt$1(reciprocal$1(cast(1, x.dtype).sub(x.ref.mul(x.ref))));
+		return [[asin$1(x)], [denom.mul(dx)]];
+	},
+	[Primitive.Atan]([x], [dx]) {
+		const denom = cast(1, x.dtype).add(x.ref.mul(x.ref));
+		return [[atan$1(x)], [dx.div(denom)]];
+	},
 	[Primitive.Exp]([x], [dx]) {
 		const z = exp$1(x);
 		return [[z.ref], [z.mul(dx)]];
@@ -3019,13 +3143,14 @@ const jvpRules = {
 		const indicesRef = indices.map((t) => t.ref);
 		return [[gather(x, indices, axis, outDim)], [gather(dx, indicesRef, axis, outDim)]];
 	},
-	[Primitive.JitCall](primals, tangents, { jaxpr }) {
+	[Primitive.JitCall](primals, tangents, { name, jaxpr }) {
 		const { newJaxpr, newConsts } = jvpJaxpr(jaxpr);
 		const outs = bind(Primitive.JitCall, [
 			...newConsts.map((c) => c.ref),
 			...primals,
 			...tangents
 		], {
+			name: `${name}_jvp`,
 			jaxpr: newJaxpr,
 			numConsts: newConsts.length
 		});
@@ -3080,12 +3205,15 @@ var import_usingCtx = /* @__PURE__ */ __toESM(require_usingCtx(), 1);
 function mappedAval(batchDim, aval) {
 	const shape$1 = [...aval.shape];
 	shape$1.splice(batchDim, 1);
-	return new ShapedArray(shape$1, aval.dtype);
+	return new ShapedArray(shape$1, aval.dtype, aval.weakType);
 }
 /** Move one axis to a different index. */
 function moveaxis$1(x, src, dst) {
 	const t = pureArray(x);
-	const perm = require_backend.range(t.shape.length);
+	src = require_backend.checkAxis(src, t.ndim);
+	dst = require_backend.checkAxis(dst, t.ndim);
+	if (src === dst) return t;
+	const perm = require_backend.range(t.ndim);
 	perm.splice(src, 1);
 	perm.splice(dst, 0, src);
 	return transpose$1(t, perm);
@@ -3178,6 +3306,8 @@ const vmapRules = {
 	[Primitive.Bitcast]: unopBatcher((x, { dtype }) => bitcast(x, dtype)),
 	[Primitive.Sin]: unopBatcher(sin$1),
 	[Primitive.Cos]: unopBatcher(cos$1),
+	[Primitive.Asin]: unopBatcher(asin$1),
+	[Primitive.Atan]: unopBatcher(atan$1),
 	[Primitive.Exp]: unopBatcher(exp$1),
 	[Primitive.Log]: unopBatcher(log$1),
 	[Primitive.Sqrt]: unopBatcher(sqrt$1),
@@ -3219,9 +3349,10 @@ const vmapRules = {
 		const newWidth = width.toSpliced(xBdim, 0, [0, 0]);
 		return [[pad$1(x, newWidth)], [xBdim]];
 	},
-	[Primitive.JitCall](axisSize, args, dims, { jaxpr }) {
+	[Primitive.JitCall](axisSize, args, dims, { name, jaxpr }) {
 		const { newJaxpr, newConsts } = vmapJaxpr(jaxpr, axisSize, dims);
 		const outs = bind(Primitive.JitCall, [...newConsts.map((c) => c.ref), ...args], {
+			name: `${name}_vmap`,
 			jaxpr: newJaxpr,
 			numConsts: newConsts.length
 		});
@@ -3237,7 +3368,7 @@ function vmapJaxpr(jaxpr, axisSize, dims) {
 		if (dims[i] === null) return v.aval;
 		const shape$1 = [...v.aval.shape];
 		shape$1.splice(dims[i], 0, axisSize);
-		return new ShapedArray(shape$1, v.aval.dtype);
+		return new ShapedArray(shape$1, v.aval.dtype, v.aval.weakType);
 	});
 	const { jaxpr: newJaxpr, consts: newConsts } = makeJaxpr$1((args) => vmapFlat(jaxprAsFun(jaxpr), dims, args))(inAvals);
 	const result = {
@@ -3363,20 +3494,28 @@ function linearizeFlatUtil(f, primalsIn) {
 function linearizeFlat(f, primalsIn) {
 	const { primalsOut, jaxpr, consts } = linearizeFlatUtil(f, primalsIn);
 	const fLin = (...tangents) => evalJaxpr(jaxpr, [...consts.map((c) => c.ref), ...tangents]);
-	return [primalsOut, fLin];
+	const dispose$1 = () => {
+		for (const c of consts) c.dispose();
+	};
+	return [
+		primalsOut,
+		fLin,
+		dispose$1
+	];
 }
 function linearize$1(f, ...primalsIn) {
 	const [primalsInFlat, inTree] = flatten(primalsIn);
 	const [fFlat, outTree] = flattenFun(f, inTree);
-	const [primalsOutFlat, fLinFlat] = linearizeFlat(fFlat, primalsInFlat.map(pureArray));
+	const [primalsOutFlat, fLinFlat, dispose$1] = linearizeFlat(fFlat, primalsInFlat.map(pureArray));
 	if (outTree.value === void 0) throw new Error("outTree was not set in linearize");
 	const primalsOut = unflatten(outTree.value, primalsOutFlat);
-	const fLin = (...tangentsIn) => {
+	const fLin = ((...tangentsIn) => {
 		const [tangentsInFlat, inTree2] = flatten(tangentsIn);
 		if (!inTree.equals(inTree2)) throw new TreeMismatchError("linearize", inTree, inTree2);
 		const tangentsOutFlat = fLinFlat(...tangentsInFlat.map(pureArray));
 		return unflatten(outTree.value, tangentsOutFlat);
-	};
+	});
+	fLin.dispose = dispose$1;
 	return [primalsOut, fLin];
 }
 var PartialEvalTracer = class extends Tracer {
@@ -3442,8 +3581,8 @@ var PartialEvalTrace = class extends Trace {
 	processPrimitive(primitive, tracers, params) {
 		if (tracers.every((t) => t.pval.isKnown)) return bind(primitive, tracers.map((t) => t.fullLower()), params);
 		if (primitive === Primitive.JitCall) {
-			const { jaxpr, numConsts } = params;
-			return this.#partialEvalJaxpr(jaxpr, numConsts, tracers);
+			const { name, jaxpr, numConsts } = params;
+			return this.#partialEvalJaxpr(name, jaxpr, numConsts, tracers);
 		}
 		const tracersIn = tracers.map((t) => this.instantiateConst(t));
 		const avalsIn = tracersIn.map((t) => t.pval.aval);
@@ -3469,12 +3608,13 @@ var PartialEvalTrace = class extends Trace {
 	*
 	* Used when encountering a JitCall rule during the trace.
 	*/
-	#partialEvalJaxpr(jaxpr, numConsts, tracers) {
+	#partialEvalJaxpr(name, jaxpr, numConsts, tracers) {
 		jaxpr = jaxpr.flatten();
 		const inUnknowns = tracers.map((t) => !t.pval.isKnown);
 		const { jaxpr1, jaxpr2, outUnknowns, numRes } = partialEvalJaxpr(jaxpr, inUnknowns);
 		const [knownTracers, unknownTracers] = require_backend.partitionList(inUnknowns, tracers);
 		const outs1Res = bind(Primitive.JitCall, knownTracers.map((t) => t.ref.fullLower()), {
+			name: `${name}_peval`,
 			jaxpr: jaxpr1,
 			numConsts: 0
 		});
@@ -3486,13 +3626,17 @@ var PartialEvalTrace = class extends Trace {
 			prim: Primitive.JitCall,
 			tracersIn: resTracers.concat(unknownTracers),
 			params: {
+				name: `${name}_resid`,
 				jaxpr: jaxpr2,
 				numConsts: 0
 			},
 			avalsOut: jaxpr2.outs.map((x) => x.aval),
 			tracerRefsOut: []
 		};
-		const outs2 = jaxpr2.outs.map((x) => new PartialEvalTracer(this, PartialVal.unknown(x.aval), recipe));
+		const outs2 = jaxpr2.outs.map((x, i$1) => {
+			if (i$1 > 0) recipe.tracersIn.forEach((t) => t.ref);
+			return new PartialEvalTracer(this, PartialVal.unknown(x.aval), recipe);
+		});
 		recipe.tracerRefsOut = outs2.map((t) => new WeakRef(t));
 		let i = 0;
 		let j = 0;
@@ -3576,13 +3720,15 @@ function partialEvalGraphToJaxpr(tracersIn, tracersOut) {
 	const [consts, constvars] = require_backend.unzip2(constToVar.entries());
 	const inBinders = [...constvars, ...tracersIn.map((t) => tracerToVar.get(t))];
 	const outVars = tracersOut.map((t) => tracerToVar.get(t));
-	const jaxpr = new Jaxpr(inBinders, eqns, outVars);
+	let jaxpr = new Jaxpr(inBinders, eqns, outVars);
 	typecheckJaxpr(jaxpr);
 	for (const t of consts) t.ref;
 	for (const t of tracersIn) t.dispose();
 	for (const t of tracersOut) t.dispose();
+	jaxpr = jaxpr.simplify();
+	if (require_backend.DEBUG >= 5) console.log("jaxpr from partial evaluation:\n" + jaxpr.toString());
 	return {
-		jaxpr: jaxpr.simplify(),
+		jaxpr,
 		consts
 	};
 }
@@ -3623,7 +3769,7 @@ function evalJaxprTransposed(jaxpr, args, cotangents) {
 	}
 	for (let i = jaxpr.eqns.length - 1; i >= 0; i--) {
 		const eqn = jaxpr.eqns[i];
-		const primalsIn = eqn.inputs.map((v) => v instanceof Lit ? scalar(v.value, { dtype: v.dtype }) : knownPrimals.has(v) ? knownPrimals.get(v).ref : new UndefPrimal(v.aval));
+		const primalsIn = eqn.inputs.map((v) => v instanceof Lit ? array(v.value, { dtype: v.dtype }) : knownPrimals.has(v) ? knownPrimals.get(v).ref : new UndefPrimal(v.aval));
 		const cotangentsOut = eqn.outBinders.map(readCotangent);
 		const rule = transposeRules[eqn.primitive];
 		if (!rule) throw new TypeError(`Backward pass not implemented for ${eqn.primitive}`);
@@ -3708,7 +3854,7 @@ const transposeRules = {
 	},
 	[Primitive.Dot]([ct], [x, y]) {
 		if (x instanceof UndefPrimal === y instanceof UndefPrimal) throw new NonlinearError(Primitive.Dot);
-		const axisSize = generalBroadcast(x.aval.shape, y.aval.shape).slice(-1)[0];
+		const axisSize = require_backend.generalBroadcast(x.aval.shape, y.aval.shape).slice(-1)[0];
 		ct = broadcast(ct, ct.shape.concat(axisSize), [-1]);
 		return [x instanceof UndefPrimal ? unbroadcast(mul(ct, y), x) : null, y instanceof UndefPrimal ? unbroadcast(mul(x, ct), y) : null];
 	},
@@ -3803,7 +3949,7 @@ const transposeRules = {
 		if (indices.some((i) => i instanceof UndefPrimal)) throw new NonlinearError(Primitive.Gather);
 		throw new Error("Gather transpose rule is not yet implemented, requires complex Scatter sum operation");
 	},
-	[Primitive.JitCall](cts, args, { jaxpr }) {
+	[Primitive.JitCall](cts, args, { name, jaxpr }) {
 		const undefPrimals = args.map((x) => x instanceof UndefPrimal);
 		const { newJaxpr, newConsts } = transposeJaxpr(jaxpr, undefPrimals);
 		const residuals = args.filter((x, i$1) => !undefPrimals[i$1]);
@@ -3812,6 +3958,7 @@ const transposeRules = {
 			...residuals,
 			...cts
 		], {
+			name: `${name}_t`,
 			jaxpr: newJaxpr,
 			numConsts: newConsts.length
 		});
@@ -3848,20 +3995,28 @@ function vjpFlat(f, primalsIn) {
 		const transposeInputs = [...consts.map((c) => c.ref), ...primalsIn.map((t) => new UndefPrimal(t.aval))];
 		return evalJaxprTransposed(jaxpr, transposeInputs, cotangents);
 	};
-	return [primalsOut, fVjp];
+	const dispose$1 = () => {
+		for (const c of consts) c.dispose();
+	};
+	return [
+		primalsOut,
+		fVjp,
+		dispose$1
+	];
 }
 function vjp$1(f, ...primalsIn) {
 	const [primalsInFlat, inTree] = flatten(primalsIn);
 	const [fFlat, outTree] = flattenFun(f, inTree);
-	const [primalsOutFlat, fVjpFlat] = vjpFlat(fFlat, primalsInFlat.map(pureArray));
+	const [primalsOutFlat, fVjpFlat, dispose$1] = vjpFlat(fFlat, primalsInFlat.map(pureArray));
 	if (outTree.value === void 0) throw new Error("outTree was not set in vjp");
 	const primalsOut = unflatten(outTree.value, primalsOutFlat);
-	const fVjp = (cotangentsOut) => {
+	const fVjp = ((cotangentsOut) => {
 		const [cotangentsOutFlat, outTree2] = flatten(cotangentsOut);
 		if (!outTree.value.equals(outTree2)) throw new TreeMismatchError("vjp", outTree.value, outTree2);
 		const cotangentsInFlat = fVjpFlat(...cotangentsOutFlat.map(pureArray));
 		return unflatten(inTree, cotangentsInFlat);
-	};
+	});
+	fVjp.dispose = dispose$1;
 	return [primalsOut, fVjp];
 }
 function grad$1(f) {
@@ -3878,8 +4033,9 @@ function valueAndGrad$1(f) {
 		const [y, fVjp] = vjp$1(f, x[0], ...x.slice(1).map(stopGradient));
 		if (!(y instanceof Tracer) || ndim$1(y) !== 0) throw new TypeError("grad requires a scalar output");
 		if (!require_backend.isFloatDtype(y.dtype)) throw new TypeError("grad only supports floating-point dtypes");
-		const [ct, ...rest] = fVjp(scalar(1, { dtype: y.dtype }));
-		for (const r of rest) r.dispose();
+		const [ct, ...rest] = fVjp(array(1, { dtype: y.dtype }));
+		for (const r of rest) dispose(r);
+		fVjp.dispose();
 		return [y, ct];
 	};
 }
@@ -3887,7 +4043,13 @@ function jacrev$1(f) {
 	return function jacobianReverse(x) {
 		if (x.shape.length !== 1) throw new TypeError("jacrev only supports 1D inputs");
 		const [size$1] = x.shape;
-		const pullback = (ct) => vjp$1(f, x)[1](ct)[0];
+		const pullback = (ct) => {
+			const [y, fVjp] = vjp$1(f, x);
+			y.dispose();
+			const [ret] = fVjp(ct);
+			fVjp.dispose();
+			return ret;
+		};
 		return vmap$1(pullback, [1])(eye(size$1, void 0, { dtype: x.dtype }));
 	};
 }
@@ -3967,19 +4129,38 @@ __export(numpy_exports, {
 	DType: () => require_backend.DType,
 	abs: () => abs,
 	absolute: () => absolute,
+	acos: () => acos,
+	acosh: () => acosh,
 	add: () => add,
 	allclose: () => allclose,
 	arange: () => arange,
+	arccos: () => arccos,
+	arccosh: () => arccosh,
+	arcsinh: () => arcsinh,
+	arctan: () => arctan,
+	arctan2: () => arctan2,
+	arctanh: () => arctanh,
 	argmax: () => argmax,
 	argmin: () => argmin,
 	array: () => array,
+	asin: () => asin,
+	asinh: () => asinh,
 	astype: () => astype,
+	atan: () => atan,
+	atan2: () => atan2,
+	atanh: () => atanh,
 	bool: () => bool,
+	broadcastArrays: () => broadcastArrays,
+	broadcastShapes: () => broadcastShapes,
+	broadcastTo: () => broadcastTo,
+	cbrt: () => cbrt,
 	clip: () => clip,
 	columnStack: () => columnStack,
 	concatenate: () => concatenate,
 	cos: () => cos,
 	cosh: () => cosh,
+	deg2rad: () => deg2rad,
+	degrees: () => degrees,
 	diag: () => diag,
 	diagonal: () => diagonal,
 	divide: () => divide,
@@ -3990,6 +4171,7 @@ __export(numpy_exports, {
 	eulerGamma: () => eulerGamma,
 	exp: () => exp,
 	exp2: () => exp2,
+	expm1: () => expm1,
 	eye: () => eye,
 	flip: () => flip,
 	fliplr: () => fliplr,
@@ -4001,14 +4183,17 @@ __export(numpy_exports, {
 	greater: () => greater,
 	greaterEqual: () => greaterEqual,
 	hstack: () => hstack,
+	hypot: () => hypot,
 	identity: () => identity$1,
 	inf: () => inf,
+	inner: () => inner,
 	int32: () => int32,
 	less: () => less,
 	lessEqual: () => lessEqual,
 	linspace: () => linspace,
 	log: () => log,
 	log10: () => log10,
+	log1p: () => log1p,
 	log2: () => log2,
 	matmul: () => matmul,
 	max: () => max,
@@ -4024,35 +4209,49 @@ __export(numpy_exports, {
 	negative: () => negative,
 	notEqual: () => notEqual,
 	ones: () => ones,
-	onesLike: () => onesLike$1,
+	onesLike: () => onesLike,
+	outer: () => outer,
 	pad: () => pad,
 	permuteDims: () => permuteDims,
 	pi: () => pi,
+	pow: () => pow,
+	power: () => power,
 	prod: () => prod$1,
+	promoteTypes: () => require_backend.promoteTypes,
+	rad2deg: () => rad2deg,
+	radians: () => radians,
 	ravel: () => ravel,
 	reciprocal: () => reciprocal,
+	repeat: () => repeat,
 	reshape: () => reshape,
-	scalar: () => scalar,
 	shape: () => shape,
+	sign: () => sign,
 	sin: () => sin,
 	sinh: () => sinh,
 	size: () => size,
 	sqrt: () => sqrt,
 	square: () => square,
 	stack: () => stack,
+	std: () => std,
+	subtract: () => subtract,
 	sum: () => sum,
 	tan: () => tan,
 	tanh: () => tanh,
+	tile: () => tile,
 	transpose: () => transpose,
+	tri: () => tri,
+	tril: () => tril,
+	triu: () => triu,
 	trueDivide: () => trueDivide,
 	trunc: () => trunc,
 	uint32: () => uint32,
+	var_: () => var_,
 	vdot: () => vdot,
 	vecdot: () => vecdot,
 	vstack: () => vstack,
 	where: () => where,
 	zeros: () => zeros,
-	zerosLike: () => zerosLike$1
+	zerosLike: () => zerosLike
 });
 const float32 = require_backend.DType.Float32;
 const int32 = require_backend.DType.Int32;
@@ -4069,54 +4268,66 @@ const inf = Number.POSITIVE_INFINITY;
 const nan = NaN;
 /** This is Pi, `π = 3.14159265358979...` */
 const pi = Math.PI;
-/** Element-wise addition, with broadcasting. */
+/** @function Element-wise addition, with broadcasting. */
 const add = add$1;
-/** Element-wise multiplication, with broadcasting. */
+/** @function Element-wise multiplication, with broadcasting. */
 const multiply = mul;
-/** Numerical negative of every element of an array. */
+/** @function Numerical negative of every element of an array. */
 const negative = neg;
-/** Calculate element-wise reciprocal of the input. This is `1/x`. */
+/** @function Calculate element-wise reciprocal of the input. This is `1/x`. */
 const reciprocal = reciprocal$1;
-/** Element-wise sine function (takes radians). */
+/** @function Element-wise sine function (takes radians). */
 const sin = sin$1;
-/** Element-wise cosine function (takes radians). */
+/** @function Element-wise cosine function (takes radians). */
 const cos = cos$1;
-/** Calculate the exponential of all elements in the input array. */
+/** @function Element-wise inverse sine function (inverse of sin). */
+const asin = asin$1;
+/** @function Element-wise inverse tangent function (inverse of tan). */
+const atan = atan$1;
+/** @function Calculate the exponential of all elements in the input array. */
 const exp = exp$1;
-/** Calculate the natural logarithm of all elements in the input array. */
+/** @function Calculate the natural logarithm of all elements in the input array. */
 const log = log$1;
-/** Calculate the square root of all elements in the input array. */
+/** @function Calculate the square root of all elements in the input array. */
 const sqrt = sqrt$1;
-/** Return element-wise minimum of the input arrays. */
+/** @function Return element-wise minimum of the input arrays. */
 const minimum = min$1;
-/** Return element-wise maximum of the input arrays. */
+/** @function Return element-wise maximum of the input arrays. */
 const maximum = max$1;
-/** Compare two arrays element-wise. */
+/** @function Compare two arrays element-wise. */
 const greater = greater$1;
-/** Compare two arrays element-wise. */
+/** @function Compare two arrays element-wise. */
 const less = less$1;
-/** Compare two arrays element-wise. */
+/** @function Compare two arrays element-wise. */
 const equal = equal$1;
-/** Compare two arrays element-wise. */
+/** @function Compare two arrays element-wise. */
 const notEqual = notEqual$1;
-/** Compare two arrays element-wise. */
+/** @function Compare two arrays element-wise. */
 const greaterEqual = greaterEqual$1;
-/** Compare two arrays element-wise. */
+/** @function Compare two arrays element-wise. */
 const lessEqual = lessEqual$1;
-/** Element-wise ternary operator, evaluates to `x` if cond else `y`. */
+/** @function Element-wise ternary operator, evaluates to `x` if cond else `y`. */
 const where = where$1;
-/** Permute the dimensions of an array. Defaults to reversing the axis order. */
+/**
+* @function
+* Permute the dimensions of an array. Defaults to reversing the axis order.
+*/
 const transpose = transpose$1;
 /**
+* @function
 * Give a new shape to an array without changing its data.
 *
 * One shape dimension can be -1. In this case, the value is inferred from the
 * length of the array and remaining dimensions.
 */
 const reshape = reshape$1;
-/** Move axes of an array to new positions. Other axes retain original order. */
+/**
+* @function
+* Move axes of an array to new positions. Other axes retain original order.
+*/
 const moveaxis = moveaxis$1;
 /**
+* @function
 * Add padding (zeros) to an array.
 *
 * The `width` argument is either an integer or pair of integers, in which case
@@ -4124,15 +4335,27 @@ const moveaxis = moveaxis$1;
 * pair specifies the padding for its corresponding axis.
 */
 const pad = pad$1;
-/** Return the number of dimensions of an array. Does not consume array reference. */
+/**
+* @function
+* Return the number of dimensions of an array. Does not consume array reference.
+*/
 const ndim = ndim$1;
-/** Return the shape of an array. Does not consume array reference. */
+/** @function Return the shape of an array. Does not consume array reference. */
 const shape = getShape;
-/** Return an array of zeros with the same shape and type as a given array. */
-const zerosLike$1 = zerosLike;
-/** Return an array of ones with the same shape and type as a given array. */
-const onesLike$1 = onesLike;
-/** Return a full array with the same shape and type as a given array. */
+/**
+* @function
+* Return an array of zeros with the same shape and type as a given array.
+*/
+const zerosLike = zerosLike$1;
+/**
+* @function
+* Return an array of ones with the same shape and type as a given array.
+*/
+const onesLike = onesLike$1;
+/**
+* @function
+* Return a full array with the same shape and type as a given array.
+*/
 const fullLike$1 = fullLike;
 /**
 * Return the number of elements in an array, optionally along an axis.
@@ -4147,23 +4370,23 @@ function astype(a, dtype) {
 	return fudgeArray(a).astype(dtype);
 }
 /** Sum of the elements of the array over a given axis, or axes. */
-function sum(a, axis, opts) {
+function sum(a, axis = null, opts) {
 	return reduce(a, require_backend.AluOp.Add, axis, opts);
 }
 /** Product of the array elements over a given axis. */
-function prod$1(a, axis, opts) {
+function prod$1(a, axis = null, opts) {
 	return reduce(a, require_backend.AluOp.Mul, axis, opts);
 }
 /** Return the minimum of array elements along a given axis. */
-function min(a, axis, opts) {
+function min(a, axis = null, opts) {
 	return reduce(a, require_backend.AluOp.Min, axis, opts);
 }
 /** Return the maximum of array elements along a given axis. */
-function max(a, axis, opts) {
+function max(a, axis = null, opts) {
 	return reduce(a, require_backend.AluOp.Max, axis, opts);
 }
 /** Compute the average of the array elements along the specified axis. */
-function mean(a, axis, opts) {
+function mean(a, axis = null, opts) {
 	return fudgeArray(a).mean(axis, opts);
 }
 /**
@@ -4179,8 +4402,8 @@ function argmin(a, axis, opts) {
 		axis = 0;
 	} else axis = require_backend.checkAxis(axis, a.ndim);
 	const shape$1 = a.shape;
-	const isMax = equal(a, min(a.ref, axis, { keepDims: true }));
-	const length = scalar(shape$1[axis], {
+	const isMax = equal(a, min(a.ref, axis, { keepdims: true }));
+	const length = array(shape$1[axis], {
 		dtype: int32,
 		device: a.device
 	});
@@ -4203,8 +4426,8 @@ function argmax(a, axis, opts) {
 		axis = 0;
 	} else axis = require_backend.checkAxis(axis, a.ndim);
 	const shape$1 = a.shape;
-	const isMax = equal(a, max(a.ref, axis, { keepDims: true }));
-	const length = scalar(shape$1[axis], {
+	const isMax = equal(a, max(a.ref, axis, { keepdims: true }));
+	const length = array(shape$1[axis], {
 		dtype: int32,
 		device: a.device
 	});
@@ -4215,17 +4438,9 @@ function argmax(a, axis, opts) {
 	return length.sub(max(idx, axis, opts));
 }
 /** Reverse the elements in an array along the given axes. */
-function flip(x, axis) {
+function flip(x, axis = null) {
 	const nd = ndim(x);
-	if (axis === void 0) axis = require_backend.range(nd);
-	else if (typeof axis === "number") axis = [axis];
-	const seen = /* @__PURE__ */ new Set();
-	for (let i = 0; i < axis.length; i++) {
-		if (axis[i] >= nd || axis[i] < -nd) throw new Error(`flip: axis ${axis[i]} out of bounds for array of ${nd} dimensions`);
-		if (axis[i] < 0) axis[i] += nd;
-		if (seen.has(axis[i])) throw new Error(`flip: duplicate axis ${axis[i]} in axis list`);
-		seen.add(axis[i]);
-	}
+	axis = require_backend.normalizeAxis(axis, nd);
 	return flip$1(x, axis);
 }
 /**
@@ -4331,12 +4546,80 @@ function flipud(x) {
 function fliplr(x) {
 	return flip(x, 1);
 }
+/** @function Alternative name for `numpy.transpose()`. */
 const permuteDims = transpose;
 /** Return a 1-D flattened array containing the elements of the input. */
 function ravel(a) {
 	return fudgeArray(a).ravel();
 }
 /**
+* Repeat each element of an array after themselves.
+*
+* If no axis is provided, use the flattened input array, and return a flat
+* output array.
+*/
+function repeat(a, repeats, axis) {
+	if (!Number.isInteger(repeats) || repeats < 0) throw new Error(`repeat: repeats must be a non-negative integer, got ${repeats}`);
+	a = fudgeArray(a);
+	if (axis === void 0) {
+		a = ravel(a);
+		axis = 0;
+	}
+	axis = require_backend.checkAxis(axis, a.ndim);
+	if (repeats === 1) return a;
+	const broadcastedShape = a.shape.toSpliced(axis + 1, 0, repeats);
+	const finalShape = a.shape.toSpliced(axis, 1, a.shape[axis] * repeats);
+	return broadcast(a, broadcastedShape, [axis + 1]).reshape(finalShape);
+}
+/**
+* Construct an array by repeating A the number of times given by reps.
+*
+* If `A` is an array of shape `(d1, d2, ..., dn)` and `reps` is a sequence of
+* integers, the resulting array will have a shape of `(reps[0] * d1,
+* reps[1] * d2, ..., reps[n] * dn)`, with `A` tiled along each dimension.
+*/
+function tile(a, reps) {
+	a = fudgeArray(a);
+	if (typeof reps === "number") reps = [reps];
+	if (!reps.every((r) => Number.isInteger(r) && r >= 0)) throw new Error(`tile: reps must be non-negative integers, got ${JSON.stringify(reps)}`);
+	const ndiff = reps.length - a.ndim;
+	if (ndiff > 0) a = a.reshape([...require_backend.rep(ndiff, 1), ...a.shape]);
+	if (ndiff < 0) reps = [...require_backend.rep(-ndiff, 1), ...reps];
+	const broadcastedShape = [];
+	const broadcastAxes = [];
+	for (let i = 0; i < a.ndim; i++) {
+		if (reps[i] > 1) {
+			broadcastedShape.push(reps[i]);
+			broadcastAxes.push(broadcastedShape.length - 1);
+		}
+		broadcastedShape.push(a.shape[i]);
+	}
+	const finalShape = a.shape.map((d, i) => reps[i] * d);
+	return broadcast(a, broadcastedShape, broadcastAxes).reshape(finalShape);
+}
+/**
+* Broadcast an array to a shape, with NumPy-style broadcasing rules.
+*
+* In other words, this lets you append axes to the left, and/or expand
+* dimensions where the shape is 1.
+*/
+function broadcastTo(a, shape$1) {
+	const nd = ndim(a);
+	if (shape$1.length < nd) throw new Error(`broadcastTo: target shape ${JSON.stringify(shape$1)} has fewer dimensions than input array: ${nd}`);
+	return broadcast(a, shape$1, require_backend.range(shape$1.length - nd));
+}
+/** Broadcast input shapes to a common output shape. */
+function broadcastShapes(...shapes) {
+	if (shapes.length === 0) return [];
+	return shapes.reduce(require_backend.generalBroadcast);
+}
+/** Broadcast arrays to a common shape. */
+function broadcastArrays(...arrays) {
+	const shapes = arrays.map((a) => shape(a));
+	const outShape = broadcastShapes(...shapes);
+	return arrays.map((a) => broadcastTo(a, outShape));
+}
+/**
 * Return specified diagonals.
 *
 * If a is 2D, return the diagonal of the array with the given offset. If a is
@@ -4360,7 +4643,7 @@ function diag(v, k = 0) {
 	if (!Number.isInteger(k)) throw new TypeError(`k must be an integer, got ${k}`);
 	if (a.ndim === 1) {
 		const n = a.shape[0];
-		const ret = where(eye(n).equal(1), a.ref, zerosLike$1(a));
+		const ret = where(eye(n).equal(1), a.ref, zerosLike(a));
 		if (k > 0) return pad(ret, [[0, k], [k, 0]]);
 		else if (k < 0) return pad(ret, [[-k, 0], [0, -k]]);
 		else return ret;
@@ -4404,8 +4687,36 @@ function dot(x, y) {
 	]);
 	return dot$1(x, y);
 }
-/** Vector dot product of two arrays. */
-function vecdot(x, y) {
+/**
+* Compute the inner product of two arrays.
+*
+* Unlike `jax.numpy.matmul()` or `jax.numpy.dot()`, this always performs a
+* contraction on the last axis.
+*
+* Returned array has shape `[...x.shape[:-1], ...y.shape[:-1]]`.
+*/
+function inner(x, y) {
+	x = reshape(x, shape(x).toSpliced(-1, 0, ...require_backend.rep(ndim(y) - 1, 1)));
+	return dot$1(x, y);
+}
+/**
+* Compute the outer product of two arrays.
+*
+* If the input arrays are not 1D, they will be flattened. Returned array will
+* be of shape `[x.size, y.size]`.
+*/
+function outer(x, y) {
+	x = ravel(x);
+	y = ravel(y);
+	return multiply(x.reshape([x.shape[0], 1]), y);
+}
+/** Vector dot product of two arrays along a given axis. */
+function vecdot(x, y, { axis } = {}) {
+	const xaxis = require_backend.checkAxis(axis ?? -1, ndim(x));
+	const yaxis = require_backend.checkAxis(axis ?? -1, ndim(y));
+	if (shape(x)[xaxis] !== shape(y)[yaxis]) throw new Error(`vecdot: shapes ${JSON.stringify(shape(x))} and ${JSON.stringify(shape(y))} not aligned along axis ${axis}: ${shape(x)[xaxis]} != ${shape(y)[yaxis]}`);
+	x = moveaxis(x, xaxis, -1);
+	y = moveaxis(y, yaxis, -1);
 	return dot$1(x, y);
 }
 /**
@@ -4414,7 +4725,7 @@ function vecdot(x, y) {
 * Like vecdot() but flattens the arguments first into vectors.
 */
 function vdot(x, y) {
-	return vecdot(ravel(x), ravel(y));
+	return dot$1(ravel(x), ravel(y));
 }
 /**
 * Return a tuple of coordinate matrices from coordinate vectors.
@@ -4443,6 +4754,43 @@ function meshgrid(xs, { indexing } = {}) {
 	return xs.map((x, i) => broadcast(x, shape$1, [...require_backend.range(i), ...require_backend.range(i + 1, xs.length)]));
 }
 /**
+* Return an array with ones on and below the diagonal and zeros elsewhere.
+*
+* If `k` is provided, it specifies the sub-diagonal on and below which the
+* array is filled with ones. `k=0` is the main diagonal, `k<0` is below it, and
+* `k>0` is above it.
+*/
+function tri(n, m, k = 0, { dtype, device } = {}) {
+	m ??= n;
+	dtype ??= require_backend.DType.Float32;
+	if (!Number.isInteger(n) || n < 0) throw new TypeError(`tri: n must be a non-negative integer, got ${n}`);
+	if (!Number.isInteger(m) || m < 0) throw new TypeError(`tri: m must be a non-negative integer, got ${m}`);
+	if (!Number.isInteger(k)) throw new TypeError(`tri: k must be an integer, got ${k}`);
+	const rows = arange(k, n + k, 1, {
+		dtype: require_backend.DType.Int32,
+		device
+	});
+	const cols = arange(0, m, 1, {
+		dtype: require_backend.DType.Int32,
+		device
+	});
+	return rows.reshape([n, 1]).greaterEqual(cols).astype(dtype);
+}
+/** Return the lower triangle of an array. Must be of dimension >= 2. */
+function tril(a, k = 0) {
+	if (ndim(a) < 2) throw new TypeError(`tril: input array must be at least 2D, got ${ndim(a)}D`);
+	a = fudgeArray(a);
+	const [n, m] = a.shape.slice(-2);
+	return where(tri(n, m, k, { dtype: bool }), a.ref, zerosLike(a));
+}
+/** Return the upper triangle of an array. Must be of dimension >= 2. */
+function triu(a, k = 0) {
+	if (ndim(a) < 2) throw new TypeError(`tril: input array must be at least 2D, got ${ndim(a)}D`);
+	a = fudgeArray(a);
+	const [n, m] = a.shape.slice(-2);
+	return where(tri(n, m, k - 1, { dtype: bool }), zerosLike(a.ref), a);
+}
+/**
 * Clip (limit) the values in an array.
 *
 * Given an interval, values outside the interval are clipped to the interval
@@ -4466,18 +4814,70 @@ function absolute(x) {
 	x = fudgeArray(x);
 	return where(less(x.ref, 0), x.ref.mul(-1), x);
 }
-/** Alias of `jax.numpy.absolute()`. */
+/** @function Alias of `jax.numpy.absolute()`. */
 const abs = absolute;
+/** Return an element-wise indication of sign of the input. */
+function sign(x) {
+	x = fudgeArray(x);
+	return where(notEqual(x.ref, 0), where(less(x.ref, 0), -1, 1), 0);
+}
 /** Calculate element-wise square of the input array. */
 function square(x) {
 	x = fudgeArray(x);
 	return x.ref.mul(x);
 }
-/** Compute a trigonometric tangent of each element of input. */
+/** Element-wise tangent function (takes radians). */
 function tan(x) {
 	x = fudgeArray(x);
 	return sin(x.ref).div(cos(x));
 }
+/** Element-wise inverse cosine function (inverse of cos). */
+function acos(x) {
+	return subtract(pi / 2, asin(x));
+}
+/**
+* @function
+* Return element-wise hypotenuse for the given legs of a right triangle.
+*
+* In the original NumPy/JAX implementation, this function is more numerically
+* stable than sqrt(x1**2 + x2**2). We don't currently implement those stability
+* improvements.
+*/
+const hypot = jit$1(function hypot$1(x1, x2) {
+	return sqrt(square(x1).add(square(x2)));
+});
+/**
+* @function
+* Element-wise arc tangent of y/x with correct quadrant.
+*
+* Returns the angle in radians between the positive x-axis and the point (x, y).
+* The result is in the range [-π, π].
+*
+* Uses numerically stable formulas:
+* - When x >= 0: atan2(y, x) = 2 * atan(y / (sqrt(x^2 + y^2) + x))
+* - When x < 0:  atan2(y, x) = 2 * atan((sqrt(x^2 + y^2) - x) / y)
+*
+* The output is ill-defined when both x and y are zero.
+*/
+const atan2 = jit$1(function atan2$1(y, x) {
+	const r = sqrt(square(x.ref).add(square(y.ref)));
+	const xNeg = less(x.ref, 0);
+	const numer = where(xNeg.ref, r.ref.sub(x.ref), y.ref);
+	const denom = where(xNeg, y, r.add(x));
+	return atan(numer.div(denom)).mul(2);
+});
+/** @function Alias of `jax.numpy.acos()`. */
+const arccos = acos;
+/** @function Alias of `jax.numpy.atan()`. */
+const arctan = atan;
+/** @function Alias of `jax.numpy.atan2()`. */
+const arctan2 = atan2;
+/** Element-wise subtraction, with broadcasting. */
+function subtract(x, y) {
+	x = fudgeArray(x);
+	y = fudgeArray(y);
+	return x.sub(y);
+}
 /** Calculates the floating-point division of x by y element-wise. */
 function trueDivide(x, y) {
 	x = fudgeArray(x);
@@ -4485,7 +4885,7 @@ function trueDivide(x, y) {
 	if (!require_backend.isFloatDtype(x.dtype) || !require_backend.isFloatDtype(y.dtype)) throw new TypeError(`trueDivide: x and y must be floating-point arrays, got ${x.dtype} and ${y.dtype}`);
 	return x.div(y);
 }
-/** Alias of `jax.numpy.trueDivide()`. */
+/** @function Alias of `jax.numpy.trueDivide()`. */
 const divide = trueDivide;
 /** Round input to the nearest integer towards zero. */
 function trunc(x) {
@@ -4503,36 +4903,134 @@ function log2(x) {
 function log10(x) {
 	return log(x).mul(Math.LOG10E);
 }
+/** Calculate `exp(x) - 1` element-wise. */
+function expm1(x) {
+	return exp(x).sub(1);
+}
+/** Calculate the natural logarithm of `1 + x` element-wise. */
+function log1p(x) {
+	return log(add(1, x));
+}
+/** Convert angles from degrees to radians. */
+function deg2rad(x) {
+	return multiply(x, pi / 180);
+}
+/** @function Alias of `jax.numpy.deg2rad()`. */
+const radians = deg2rad;
+/** Convert angles from radians to degrees. */
+function rad2deg(x) {
+	return multiply(x, 180 / pi);
+}
+/** @function Alias of `jax.numpy.rad2deg()`. */
+const degrees = rad2deg;
+/**
+* @function
+* Computes first array raised to power of second array, element-wise.
+*/
+const power = jit$1(function power$1(x1, x2) {
+	return exp(log(x1).mul(x2));
+});
+/** @function Alias of `jax.numpy.power()`. */
+const pow = power;
+/** @function Calculate the element-wise cube root of the input array. */
+const cbrt = jit$1(function cbrt$1(x) {
+	const sgn = where(less(x.ref, 0), -1, 1);
+	return sgn.ref.mul(exp(log(x.mul(sgn)).mul(1 / 3)));
+});
 /**
+* @function
 * Calculate element-wise hyperbolic sine of input.
 *
 * `sinh(x) = (exp(x) - exp(-x)) / 2`
 */
-function sinh(x) {
+const sinh = jit$1(function sinh$1(x) {
 	const ex = exp(x);
 	const emx = reciprocal(ex.ref);
 	return ex.sub(emx).mul(.5);
-}
+});
 /**
+* @function
 * Calculate element-wise hyperbolic cosine of input.
 *
 * `cosh(x) = (exp(x) + exp(-x)) / 2`
 */
-function cosh(x) {
+const cosh = jit$1(function cosh$1(x) {
 	const ex = exp(x);
 	const emx = reciprocal(ex.ref);
 	return ex.add(emx).mul(.5);
-}
+});
 /**
+* @function
 * Calculate element-wise hyperbolic tangent of input.
 *
 * `tanh(x) = sinh(x)/cosh(x) = (exp(x) - exp(-x)) / (exp(x) + exp(-x))`
 */
-function tanh(x) {
-	x = fudgeArray(x);
+const tanh = jit$1(function tanh$1(x) {
 	const negsgn = where(less(x.ref, 0), 1, -1);
 	const en2x = exp(x.mul(negsgn.ref).mul(2));
 	return en2x.ref.sub(1).div(en2x.add(1)).mul(negsgn);
+});
+/**
+* @function
+* Calculate element-wise inverse hyperbolic sine of input.
+*
+* `arcsinh(x) = ln(x + sqrt(x^2 + 1))`
+*/
+const arcsinh = jit$1(function arcsinh$1(x) {
+	return log(x.ref.add(sqrt(square(x).add(1))));
+});
+/**
+* @function
+* Calculate element-wise inverse hyperbolic cosine of input.
+*
+* `arccosh(x) = ln(x + sqrt(x^2 - 1))`
+*/
+const arccosh = jit$1(function arccosh$1(x) {
+	return log(x.ref.add(sqrt(square(x).sub(1))));
+});
+/**
+* @function
+* Calculate element-wise inverse hyperbolic tangent of input.
+*
+* `arctanh(x) = 0.5 * ln((1 + x) / (1 - x))`
+*/
+const arctanh = jit$1(function arctanh$1(x) {
+	return log(add(1, x.ref).div(subtract(1, x))).mul(.5);
+});
+/** @function Alias of `jax.numpy.arcsinh()`. */
+const asinh = arcsinh;
+/** @function Alias of `jax.numpy.arccosh()`. */
+const acosh = arccosh;
+/** @function Alias of `jax.numpy.arctanh()`. */
+const atanh = arctanh;
+/**
+* Compute the variance of an array.
+*
+* The variance is computed for the flattened array by default, otherwise over
+* the specified axis.
+*
+* If `correction` is provided, the divisor in calculation is `N - correction`,
+* where `N` represents the number of elements (e.g., for Bessel's correction).
+*/
+function var_(x, axis = null, opts) {
+	x = fudgeArray(x);
+	axis = require_backend.normalizeAxis(axis, x.ndim);
+	const n = axis.reduce((acc, a) => acc * x.shape[a], 1);
+	if (n === 0) throw new Error("var: cannot compute variance over zero-length axis");
+	const mu = opts?.mean !== void 0 ? opts.mean : mean(x.ref, axis, { keepdims: true });
+	return square(x.sub(mu)).sum(axis, { keepdims: opts?.keepdims }).mul(1 / (n - (opts?.correction ?? 0)));
+}
+/**
+* Compute the standard deviation of an array.
+*
+* The standard deviation is computed for the flattened array by default,
+* otherwise over the specified axis.
+*
+* If `correction` is provided, the divisor in calculation is `N - correction`,
+* where `N` represents the number of elements (e.g., for Bessel's correction).
+*/
+function std(x, axis = null, opts) {
+	return sqrt(var_(x, axis, opts));
 }
 
 //#endregion
@@ -4547,6 +5045,7 @@ __export(nn_exports, {
 	leakyRelu: () => leakyRelu,
 	logSigmoid: () => logSigmoid,
 	logSoftmax: () => logSoftmax,
+	logmeanexp: () => logmeanexp,
 	logsumexp: () => logsumexp,
 	mish: () => mish,
 	oneHot: () => oneHot,
@@ -4557,6 +5056,8 @@ __export(nn_exports, {
 	softSign: () => softSign,
 	softmax: () => softmax,
 	softplus: () => softplus,
+	squareplus: () => squareplus,
+	standardize: () => standardize,
 	swish: () => swish
 });
 /**
@@ -4600,6 +5101,7 @@ function softSign(x) {
 	return x.ref.div(absolute(x).add(1));
 }
 /**
+* @function
 * Sigmoid-weighted Linear Unit (SiLU) activation function, also known as
 * Swish, computed element-wise:
 * `silu(x) = x * sigmoid(x) = x / (1 + exp(-x))`.
@@ -4608,8 +5110,11 @@ function softSign(x) {
 *
 * Reference: https://en.wikipedia.org/wiki/Swish_function
 */
-const silu = jit$1((x) => x.ref.mul(sigmoid(x)));
+const silu = jit$1(function silu$1(x) {
+	return x.ref.mul(sigmoid(x));
+});
 /**
+* @function
 * Sigmoid-weighted Linear Unit (SiLU) activation function, also known as
 * Swish, computed element-wise:
 * `silu(x) = x * sigmoid(x) = x / (1 + exp(-x))`.
@@ -4626,7 +5131,10 @@ const swish = silu;
 function logSigmoid(x) {
 	return negative(softplus(negative(x)));
 }
-/** Identity activation function. Returns the argument unmodified. */
+/**
+* @function
+* Identity activation function. Returns the argument unmodified.
+*/
 const identity = fudgeArray;
 /** Leaky rectified linear (ReLU) activation function */
 function leakyRelu(x, negativeSlope = .01) {
@@ -4654,6 +5162,7 @@ function celu(x, alpha = 1) {
 	return where(less(x.ref, 0), exp(x.ref.div(alpha)).sub(1).mul(alpha), x);
 }
 /**
+* @function
 * Gaussion error linear unit (GELU) activation function.
 *
 * This is computed element-wise. Currently jax-js does not support the erf() or
@@ -4664,7 +5173,7 @@ function celu(x, alpha = 1) {
 *
 * This will be improved in the future.
 */
-const gelu = jit$1((x) => {
+const gelu = jit$1(function gelu$1(x) {
 	const SQRT_2_OVER_PI = Math.sqrt(2 / Math.PI);
 	return x.ref.mul(.5).mul(tanh(x.ref.mul(x.ref.mul(x).mul(.044715).add(1)).mul(SQRT_2_OVER_PI)).add(1));
 });
@@ -4685,6 +5194,16 @@ function glu(x, axis = -1) {
 	return a.mul(sigmoid(b));
 }
 /**
+* Squareplus activation function.
+*
+* Computes the element-wise function:
+* `squareplus(x) = 0.5 * (x + sqrt(x^2 + b))`
+*/
+function squareplus(x, b = 4) {
+	x = fudgeArray(x);
+	return x.ref.add(sqrt(square(x).add(b))).mul(.5);
+}
+/**
 * Mish activation function.
 *
 * Computes the element-wise function:
@@ -4702,17 +5221,13 @@ function mish(x) {
 *
 * Reference: https://en.wikipedia.org/wiki/Softmax_function
 */
-function softmax(x, axis) {
+function softmax(x, axis = -1) {
 	x = fudgeArray(x);
-	if (axis === void 0) axis = x.ndim ? [x.ndim - 1] : [];
-	else if (typeof axis === "number") axis = [axis];
-	if (axis.length === 0) {
-		x.dispose();
-		return ones(x.shape);
-	}
-	const xMax = max(x.ref, axis, { keepDims: true });
+	axis = require_backend.normalizeAxis(axis, x.ndim);
+	if (axis.length === 0) return onesLike(x);
+	const xMax = max(x.ref, axis, { keepdims: true });
 	const unnormalized = exp(x.sub(stopGradient(xMax)));
-	return unnormalized.ref.div(unnormalized.sum(axis, { keepDims: true }));
+	return unnormalized.ref.div(unnormalized.sum(axis, { keepdims: true }));
 }
 /**
 * Log-Softmax function.
@@ -4722,17 +5237,13 @@ function softmax(x, axis) {
 *
 * If `axis` is not specified, it defaults to the last axis.
 */
-function logSoftmax(x, axis) {
+function logSoftmax(x, axis = -1) {
 	x = fudgeArray(x);
-	if (axis === void 0) axis = x.ndim ? [x.ndim - 1] : [];
-	else if (typeof axis === "number") axis = [axis];
-	if (axis.length === 0) {
-		x.dispose();
-		return zeros(x.shape);
-	}
-	const xMax = max(x.ref, axis, { keepDims: true });
+	axis = require_backend.normalizeAxis(axis, x.ndim);
+	if (axis.length === 0) return zerosLike(x);
+	const xMax = max(x.ref, axis, { keepdims: true });
 	const shifted = x.sub(stopGradient(xMax));
-	const shiftedLogsumexp = log(exp(shifted.ref).sum(axis, { keepDims: true }));
+	const shiftedLogsumexp = log(exp(shifted.ref).sum(axis, { keepdims: true }));
 	return shifted.sub(shiftedLogsumexp);
 }
 /**
@@ -4743,16 +5254,39 @@ function logSoftmax(x, axis) {
 *
 * Reference: https://en.wikipedia.org/wiki/LogSumExp
 */
-function logsumexp(x, axis) {
+function logsumexp(x, axis = null) {
 	x = fudgeArray(x);
-	if (axis === void 0) axis = require_backend.range(x.ndim);
-	else if (typeof axis === "number") axis = [axis];
+	axis = require_backend.normalizeAxis(axis, x.ndim);
 	if (axis.length === 0) return x;
 	const xMax = stopGradient(max(x.ref, axis));
 	const xMaxDims = broadcast(xMax.ref, x.shape, axis);
 	const shifted = x.sub(xMaxDims);
 	return xMax.add(log(exp(shifted).sum(axis)));
 }
+/** Log-mean-exp reduction, like `jax.nn.logsumexp()` but subtracts `log(n)`. */
+function logmeanexp(x, axis = null) {
+	x = fudgeArray(x);
+	axis = require_backend.normalizeAxis(axis, x.ndim);
+	if (axis.length === 0) return x;
+	const n = axis.reduce((acc, a) => acc * x.shape[a], 1);
+	return logsumexp(x, axis).sub(Math.log(n));
+}
+/**
+* Standardizes input to zero mean and unit variance.
+*
+* By default, this is computed over the last axis. You can pass in a different
+* axis, or `null` to standardize over all elements.
+*
+* Epsilon is added to denominator, it defaults to `1e-5` for stability.
+*/
+function standardize(x, axis = -1, opts = {}) {
+	x = fudgeArray(x);
+	axis = require_backend.normalizeAxis(axis, x.ndim);
+	if (axis.length === 0) return x;
+	const mu = opts.mean !== void 0 ? fudgeArray(opts.mean) : x.ref.mean(axis, { keepdims: true });
+	const sigma2 = opts.variance !== void 0 ? fudgeArray(opts.variance) : square(x.ref).mean(axis, { keepdims: true }).sub(square(mu.ref));
+	return x.sub(mu).div(sqrt(sigma2.add(opts.epsilon ?? 1e-5)));
+}
 /**
 * One-hot encodes the given indices.
 *
@@ -4770,7 +5304,7 @@ function logsumexp(x, axis) {
 * ```
 */
 function oneHot(x, numClasses) {
-	if (x.dtype !== require_backend.DType.Int32) throw new TypeError(`oneHot expects integers, got ${x.dtype}`);
+	if (require_backend.isFloatDtype(x.dtype)) throw new TypeError(`oneHot expects integers, got ${x.dtype}`);
 	return eye(numClasses, void 0, { device: x.device }).slice(x);
 }
 
@@ -4778,8 +5312,11 @@ function oneHot(x, numClasses) {
 //#region src/random.ts
 var random_exports = {};
 __export(random_exports, {
+	bernoulli: () => bernoulli,
 	bits: () => bits,
+	exponential: () => exponential,
 	key: () => key,
+	normal: () => normal,
 	split: () => split,
 	uniform: () => uniform
 });
@@ -4807,21 +5344,58 @@ function bits(key$1, shape$1 = []) {
 	const keyShape = validateKeyShape(key$1);
 	return randomBits(key$1.ref.slice(...keyShape.map(() => null), 0), key$1.slice(...keyShape.map(() => null), 1), shape$1);
 }
-/** Sample uniform random values in [minval, maxval) with given shape. */
-function uniform(key$1, shape$1 = [], { minval = 0, maxval = 1 } = {}) {
+/**
+* @function
+* Sample uniform random values in [minval, maxval) with given shape.
+*/
+const uniform = jit$1(function uniform$1(key$1, shape$1 = [], { minval = 0, maxval = 1 } = {}) {
 	if (minval >= maxval) throw new Error(`Invalid range: [${minval}, ${maxval}).`);
-	const mantissa = bits(key$1, shape$1).div(scalar(512, {
+	const mantissa = bits(key$1, shape$1).div(array(512, {
 		dtype: require_backend.DType.Uint32,
 		device: key$1.device
 	}));
-	const float12 = mantissa.add(scalar(1065353216, {
+	const float12 = mantissa.add(array(1065353216, {
 		dtype: require_backend.DType.Uint32,
 		device: key$1.device
 	}));
 	const rand = bitcast(float12, require_backend.DType.Float32).sub(1);
 	if (minval === 0 && maxval === 1) return rand;
 	else return rand.mul(maxval - minval).add(minval);
+}, { staticArgnums: [1, 2] });
+/**
+* Sample Bernoulli random variables with given mean (0,1 categorical).
+*
+* Returns a random Boolean array with the specified shape. `p` can be an array
+* and must be broadcastable to `shape`.
+*/
+function bernoulli(key$1, p = .5, shape$1 = []) {
+	p = fudgeArray(p);
+	return uniform(key$1, shape$1).less(p);
 }
+/**
+* @function
+* Sample exponential random values according to `p(x) = exp(-x)`.
+*/
+const exponential = jit$1(function exponential$1(key$1, shape$1 = []) {
+	const u = uniform(key$1, shape$1);
+	return negative(log1p(negative(u)));
+}, { staticArgnums: [1] });
+/**
+* @function
+* Sample random values according to `p(x) = 1/sqrt(2pi) * exp(-x^2/2)`.
+*
+* Unlike JAX, this uses the Box-Muller transform. JAX uses the erf_inv primitive instead and
+* directly inverts the CDF, but we don't have support for that yet. Outputs will not be
+* bitwise identical to JAX.
+*/
+const normal = jit$1(function normal$1(key$1, shape$1 = []) {
+	const [k1, k2] = split(key$1, 2);
+	const u1 = uniform(k1, shape$1);
+	const u2 = uniform(k2, shape$1);
+	const radius = sqrt(log1p(negative(u1)).mul(-2));
+	const theta = u2.mul(2 * Math.PI);
+	return radius.mul(cos(theta));
+}, { staticArgnums: [1] });
 
 //#endregion
 //#region src/polyfills.ts
@@ -4831,20 +5405,36 @@ Symbol.asyncDispose ??= Symbol.for("Symbol.asyncDispose");
 
 //#endregion
 //#region src/index.ts
-/** Compute the forward-mode Jacobian-vector product for a function. */
+/**
+* @function
+* Compute the forward-mode Jacobian-vector product for a function.
+*/
 const jvp = jvp$1;
-/** Vectorize an operation on a batched axis for one or more inputs. */
+/**
+* @function
+* Vectorize an operation on a batched axis for one or more inputs.
+*/
 const vmap = vmap$1;
-/** Compute the Jacobian evaluated column-by-column by forward-mode AD. */
+/**
+* @function
+* Compute the Jacobian evaluated column-by-column by forward-mode AD.
+*/
 const jacfwd = jacfwd$1;
-/** Construct a Jaxpr by dynamically tracing a function with example inputs. */
+/**
+* @function
+* Construct a Jaxpr by dynamically tracing a function with example inputs.
+*/
 const makeJaxpr = makeJaxpr$1;
 /**
+* @function
 * Mark a function for automatic JIT compilation, with operator fusion.
 *
 * The function will be compiled the first time it is called with a set of
 * argument shapes.
 *
+* You can call `.dispose()` on the returned, JIT-compiled function after all
+* calls to free memory associated with array constants.
+*
 * **Options:**
 * - `staticArgnums`: An array of argument indices to treat as static
 *   (compile-time constant). These arguments must be hashable, won't be traced,
@@ -4854,26 +5444,59 @@ const makeJaxpr = makeJaxpr$1;
 */
 const jit = jit$1;
 /**
+* @function
 * Produce a local linear approximation to a function at a point using jvp() and
 * partial evaluation.
 */
 const linearize = linearize$1;
-/** Calculate the reverse-mode vector-Jacobian product for a function. */
+/**
+* @function
+* Calculate the reverse-mode vector-Jacobian product for a function.
+*/
 const vjp = vjp$1;
 /**
+* @function
 * Compute the gradient of a scalar-valued function `f` with respect to its
 * first argument.
 */
 const grad = grad$1;
-/** Create a function that evaluates both `f` and the gradient of `f`. */
+/**
+* @function
+* Create a function that evaluates both `f` and the gradient of `f`.
+*/
 const valueAndGrad = valueAndGrad$1;
-/** Compute the Jacobian evaluated row-by-row by reverse-mode AD. */
+/**
+* @function
+* Compute the Jacobian evaluated row-by-row by reverse-mode AD.
+*/
 const jacrev = jacrev$1;
-/** Compute the Jacobian with reverse-mode AD. Alias for `jacrev()`. */
+/**
+* @function
+* Compute the Jacobian with reverse-mode AD. Alias for `jacrev()`.
+*/
 const jacobian = jacrev;
+/**
+* Wait until all `Array` leaves are ready by calling `Array.blockUntilReady()`.
+*
+* This can be used to wait for the results of an intermediate computation to
+* finish. It's recommended to call this regularly in an iterative computation
+* to avoid queueing up too many pending operations.
+*
+* Does not consume reference to the arrays.
+*/
+async function blockUntilReady(x) {
+	const promises = [];
+	for (const leaf of leaves(x)) if (leaf instanceof Array$1) promises.push(leaf.blockUntilReady());
+	await Promise.all(promises);
+	return x;
+}
 
 //#endregion
+exports.Array = Array$1;
 exports.DType = require_backend.DType;
+exports.Jaxpr = Jaxpr;
+exports.blockUntilReady = blockUntilReady;
+exports.defaultDevice = require_backend.defaultDevice;
 exports.devices = require_backend.devices;
 exports.grad = grad;
 exports.init = require_backend.init;
@@ -4908,7 +5531,7 @@ Object.defineProperty(exports, 'random', {
     return random_exports;
   }
 });
-exports.setDevice = require_backend.setDevice;
+exports.setDebug = require_backend.setDebug;
 Object.defineProperty(exports, 'tree', {
   enumerable: true,
   get: function () {