@jax-js/jax 0.0.4 → 0.1.0

package/dist/index.js

@@ -1,5 +1,5 @@
 import { __export } from "./chunk-Cl8Af3a2.js";
-import { AluExp, AluGroup, AluOp, AluVar, DEBUG, DType, FpHash, Kernel, PPrint, Reduction, ShapeTracker, accessorAluExp, accessorGlobal, byteWidth, checkAxis, deepEqual, defaultDevice, devices, dtypedArray, dtypedJsArray, getBackend, init, invertPermutation, isFloatDtype, isNumberPair, isPermutation, normalizeAxis, partitionList, prod, promoteTypes, range, recursiveFlatten, rep, runWithCache, setDebug, toposort, unravelAlu, unzip2, zip, zipn } from "./backend-EBRGmEYw.js";
+import { AluExp, AluGroup, AluOp, AluVar, DEBUG, DType, FpHash, Kernel, PPrint, Reduction, ShapeTracker, accessorAluExp, accessorGlobal, assertNonNull, byteWidth, checkAxis, deepEqual, defaultDevice, devices, dtypedArray, dtypedJsArray, generalBroadcast, getBackend, init, invertPermutation, isFloatDtype, isNumberPair, isPermutation, normalizeAxis, partitionList, prod, promoteTypes, range, recursiveFlatten, rep, runWithCache, setDebug, toposort, unravelAlu, unzip2, zip, zipn } from "./backend-DwIAd0AG.js";
 
 //#region src/tree.ts
 var tree_exports = {};
@@ -29,6 +29,10 @@ var JsTreeDef = class JsTreeDef {
 		this.nodeMetadata = nodeMetadata;
 		this.childTreedefs = childTreedefs;
 	}
+	/** Get the total number of leaves in the tree. */
+	get size() {
+		return this.nodeType === NodeType.Leaf ? 1 : this.childTreedefs.reduce((a, b) => a + b.size, 0);
+	}
 	/** Returns a string representation of this tree definition. */
 	toString(root = true) {
 		if (root) return "JsTreeDef(" + this.toString(false) + ")";
@@ -184,6 +188,16 @@ function pool(st, ks, strides = 1, dilation = 1) {
 	const s_ = strides;
 	const d_ = dilation;
 	const o_ = zipn(i_, d_, ks, s_).map(([i, d, k, s]) => Math.ceil((i - d * (k - 1)) / s));
+	if (d_.every((d) => d === 1) && ks.every((k, j) => k <= s_[j])) {
+		st = st.padOrShrink([...noop.map(() => [0, 0]), ...zipn(i_, o_, s_).map(([i, o, s]) => [0, o * s - i])]);
+		st = st.reshape([...noop, ...zip(o_, s_).flatMap(([o, s]) => [o, s])]).shrink([...noop.map((x) => [0, x]), ...zip(o_, ks).flatMap(([o, k]) => [[0, o], [0, k]])]);
+		st = st.permute([
+			...range(noop.length),
+			...ks.map((_, j) => noop.length + 2 * j),
+			...ks.map((_, j) => noop.length + 2 * j + 1)
+		]);
+		return st;
+	}
 	const f_ = zipn(o_, s_, i_, d_, ks).map(([o, s, i, d, k]) => 1 + Number(o * s > i - d * (k - 1)));
 	const kidf = zipn(ks, i_, d_, f_);
 	st = st.repeat([...rep(noop.length, 1), ...kidf.map(([k, i, d, f]) => Math.ceil(k * (i * f + d) / i))]);
@@ -218,6 +232,12 @@ function poolTranspose(st, inShape, ks, strides = 1, dilation = 1) {
 	const s_ = strides;
 	const d_ = dilation;
 	const o_ = zipn(i_, d_, ks, s_).map(([i, d, k, s]) => Math.ceil((i - d * (k - 1)) / s));
+	if (d_.every((d) => d === 1) && ks.every((k, j) => k <= s_[j])) {
+		st = st.permute([...range(noop.length), ...ks.flatMap((_, j) => [noop.length + j, noop.length + o_.length + j])]);
+		st = st.pad([...noop.map(() => [0, 0]), ...zip(s_, ks).flatMap(([s, k]) => [[0, 0], [0, s - k]])]).reshape([...noop, ...zip(o_, s_).map(([o, s]) => o * s)]);
+		st = st.padOrShrink([...noop.map(() => [0, 0]), ...zipn(i_, o_, s_).map(([i, o, s]) => [0, i - o * s])]);
+		return st.reshape(st.shape.concat(rep(ks.length, 1)));
+	}
 	if (!deepEqual(o_, st.shape.slice(noop.length, noop.length + ks.length))) throw new Error("poolTranspose() called with mismatched output shape");
 	const f_ = zipn(o_, s_, i_, d_, ks).map(([o, s, i, d, k]) => 1 + Number(o * s > i - d * (k - 1)));
 	const kidf = zipn(ks, i_, d_, f_);
@@ -327,6 +347,8 @@ let Primitive = /* @__PURE__ */ function(Primitive$1) {
 	Primitive$1["Atan"] = "atan";
 	Primitive$1["Exp"] = "exp";
 	Primitive$1["Log"] = "log";
+	Primitive$1["Erf"] = "erf";
+	Primitive$1["Erfc"] = "erfc";
 	Primitive$1["Sqrt"] = "sqrt";
 	Primitive$1["Min"] = "min";
 	Primitive$1["Max"] = "max";
@@ -404,6 +426,12 @@ function exp$1(x) {
 function log$1(x) {
 	return bind1(Primitive.Log, [x]);
 }
+function erf$1(x) {
+	return bind1(Primitive.Erf, [x]);
+}
+function erfc$1(x) {
+	return bind1(Primitive.Erfc, [x]);
+}
 function sqrt$1(x) {
 	return bind1(Primitive.Sqrt, [x]);
 }
@@ -565,6 +593,21 @@ 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 = generalBroadcast(a.shape, b.shape);
+	const weakType = a.weakType && b.weakType;
+	let dtype;
+	if (a.weakType === b.weakType) dtype = promoteTypes(a.dtype, b.dtype);
+	else if (a.weakType) dtype = promoteTypes(b.dtype, DType.Uint32);
+	else dtype = promoteTypes(a.dtype, DType.Uint32);
+	return new ShapedArray(shape$1, dtype, weakType);
+}
 var Tracer = class Tracer {
 	/** @ignore */
 	_trace;
@@ -579,10 +622,19 @@ var Tracer = class Tracer {
 	get size() {
 		return prod(this.shape);
 	}
-	/** The dtype of the array. */
+	/** 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;
@@ -819,12 +871,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;
@@ -838,7 +891,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" ? DType.Bool : DType.Float32);
+	else if (typeof x === "boolean" || typeof x === "number") return new ShapedArray([], typeof x === "boolean" ? DType.Bool : DType.Float32, typeof x === "boolean" ? false : true);
 	else throw new TypeError(`Unknown value: ${x}`);
 }
 function bind(prim, args, params = {}) {
@@ -1121,12 +1174,18 @@ function reshapeViews(exp$2, mapping, reduceAxis = false) {
 		} else if (exp$3.op === AluOp.GlobalIndex) throw new Error("internal: reshapeViews() called with GlobalIndex op");
 	});
 }
-function broadcastedJit(fn) {
+function broadcastedJit(fn, opts) {
 	return (nargs, exps, avals, params) => {
-		const newShape = avals.map((aval) => aval.shape).reduce(generalBroadcast);
-		exps = exps.map((exp$3) => reshapeViews(exp$3, (st) => {
-			if (!deepEqual(st.shape, newShape)) return st.broadcast(newShape, range(newShape.length - st.shape.length));
-		}));
+		let { shape: newShape, dtype: newDtype } = avals.reduce(promoteAvals);
+		const skipCastIdx = opts?.skipCastIdx ?? [];
+		if (skipCastIdx.length) newDtype = avals.filter((_, i) => !skipCastIdx.includes(i)).reduce(promoteAvals).dtype;
+		exps = exps.map((exp$3, i) => {
+			exp$3 = reshapeViews(exp$3, (st) => {
+				if (!deepEqual(st.shape, newShape)) return st.broadcast(newShape, range(newShape.length - st.shape.length));
+			});
+			if (exp$3.dtype !== newDtype && !skipCastIdx.includes(i)) exp$3 = AluExp.cast(newDtype, exp$3);
+			return exp$3;
+		});
 		const exp$2 = fn(exps, params);
 		return new Kernel(nargs, prod(newShape), exp$2);
 	};
@@ -1160,7 +1219,7 @@ const jitRules = {
 		const k1 = reshapeViews(keys[1], mapping);
 		const c0 = AluExp.u32(0);
 		const c1 = AluExp.cast(DType.Uint32, AluVar.gidx);
-		const exp$2 = AluExp.threefry2x32(c0, c1, k0, k1, mode);
+		const exp$2 = AluExp.threefry2x32(k0, k1, c0, c1, mode);
 		return new Kernel(nargs, prod(shape$1), exp$2);
 	},
 	[Primitive.Sin]: unopJit(AluExp.sin),
@@ -1169,6 +1228,8 @@ const jitRules = {
 	[Primitive.Atan]: unopJit(AluExp.atan),
 	[Primitive.Exp]: unopJit(AluExp.exp),
 	[Primitive.Log]: unopJit(AluExp.log),
+	[Primitive.Erf]: unopJit(AluExp.erf),
+	[Primitive.Erfc]: unopJit(AluExp.erfc),
 	[Primitive.Sqrt]: unopJit(AluExp.sqrt),
 	[Primitive.Min]: broadcastedJit(([a, b]) => AluExp.min(a, b)),
 	[Primitive.Max]: broadcastedJit(([a, b]) => AluExp.max(a, b)),
@@ -1201,7 +1262,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: AluOp.Add,
 			axis: [cs.ndim - 1]
@@ -1211,12 +1272,12 @@ const jitRules = {
 		const [stX, stY] = prepareConv(ShapeTracker.fromShape(as.shape), 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)),
-	[Primitive.Where]: broadcastedJit(([cond, a, b]) => AluExp.where(cond, a, b)),
+	[Primitive.Where]: broadcastedJit(([cond, a, b]) => AluExp.where(cond, a, b), { skipCastIdx: [0] }),
 	[Primitive.Transpose]: reshapeJit((st, { perm }) => st.permute(perm)),
 	[Primitive.Broadcast]: reshapeJit((st, { shape: shape$1, axis }) => st.broadcast(shape$1, axis)),
 	[Primitive.Reshape]: reshapeJit((st, { shape: shape$1 }) => st.reshape(shape$1)),
@@ -1265,9 +1326,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);
@@ -1386,7 +1448,7 @@ var PendingExecute = class {
 /**
 * A multidimensional numeric array with data stored on CPU or GPU.
 *
-* This is the library's core data type. Equivalent to `jnp.Array` from JAX, or
+* This is the library's core data type. Equivalent to `jax.Array` from JAX, or
 * `torch.Tensor`.
 *
 * Not to be confused with the JavaScript "Array" constructor. Avoid importing
@@ -1397,9 +1459,11 @@ var Array$1 = class Array$1 extends Tracer {
 	static #nextId = 1001;
 	id;
 	#dtype;
+	#weakType;
 	#source;
 	#st;
 	#backend;
+	#committed;
 	#rc;
 	#pendingSet;
 	/**
@@ -1408,21 +1472,23 @@ 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.#committed = args.committed;
 		this.#rc = 1;
-		this.#pendingSet = new Set(pending);
+		this.#pendingSet = new Set(args.pending);
 		if (this.#pendingSet.size === 0) this.#pendingSet = null;
-		else if (source instanceof AluExp) throw new Error("internal: AluExp source cannot have pending executes");
+		else if (this.#source instanceof 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() {
@@ -1434,6 +1500,18 @@ 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,
+			committed: args.committed ?? this.#committed,
+			pending: args.pending ?? this.#pending ?? void 0
+		});
+	}
 	get ref() {
 		this.#check();
 		this.#rc++;
@@ -1473,7 +1551,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;
 	}
@@ -1483,9 +1564,10 @@ var Array$1 = class Array$1 extends Tracer {
 	*/
 	#gather(indices, axis, outDim) {
 		this.#check();
-		if (indices.some((a) => a.#backend !== this.#backend)) throw new TypeError(`Gather indices must have the same backend: ${this.#backend.type}`);
 		const axisSet = new Set(axis);
 		if (axisSet.size !== axis.length) throw new TypeError("Gather axis must not have duplicates");
+		if (indices.some((a) => a.#committed && a.#backend !== this.#backend)) throw new TypeError(`Gather indices must have the same backend: ${this.#backend.type}`);
+		indices = indices.map((ar) => ar._putSync(this.#backend));
 		indices = Array$1.#broadcastArrays(indices);
 		const indexShape = indices[0].shape;
 		const finalShape = this.shape.filter((_, i) => !axisSet.has(i));
@@ -1522,7 +1604,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, ShapeTracker.fromShape(finalShape), this.#dtype, this.#backend, { pending });
+		return this.#newArrayFrom({
+			source: output,
+			st: ShapeTracker.fromShape(finalShape),
+			pending
+		});
 	}
 	/** Move axes to the rightmost dimension of the shape. */
 	#moveAxesDown(axis) {
@@ -1545,11 +1631,17 @@ 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 AluExp) {
 			const exp$3 = new AluExp(op, dtypeOutput, [this.#source]);
-			return new Array$1(exp$3.simplify(), this.#st, dtypeOutput, this.#backend);
+			this.dispose();
+			return this.#newArrayFrom({
+				source: exp$3.simplify(),
+				dtype: dtypeOutput,
+				weakType
+			});
 		}
 		const indices = unravelAlu(this.#st.shape, AluVar.gidx);
 		const exp$2 = new AluExp(op, dtypeOutput, [AluExp.globalView(this.#dtype, 0, this.#st, indices)]);
@@ -1559,41 +1651,67 @@ 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, ShapeTracker.fromShape(this.shape), dtypeOutput, this.#backend, { pending });
+		return this.#newArrayFrom({
+			source: output,
+			st: ShapeTracker.fromShape(this.shape),
+			dtype: dtypeOutput,
+			weakType,
+			pending
+		});
 	}
 	#binary(op, other) {
-		const custom = (src) => new AluExp(op, this.#dtype, src);
+		const custom = (src) => new 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;
-		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}`);
-			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");
+		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 (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)));
+		const weakType = castWeakType && !strongTypeOutput;
+		const { backend, committed } = Array$1.#computeBackend(name, arrays);
+		arrays = arrays.map((ar) => ar._putSync(backend));
 		arrays = Array$1.#broadcastArrays(arrays);
 		const newShape = [...arrays[0].shape];
 		if (arrays.every((ar) => ar.#source instanceof AluExp) && !reduceAxis) {
+			const sources = arrays.map((ar, i) => {
+				if (!dtypeOverride?.[i]) return AluExp.cast(castDtype, ar.#source);
+				else return ar.#source;
+			});
 			if (arrays.every((ar) => 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);
+				arrays.forEach((ar) => ar.dispose());
+				return new Array$1({
+					source: exp$4.simplify(),
+					st: arrays[0].#st,
+					dtype: exp$4.dtype,
+					weakType,
+					backend,
+					committed
+				});
 			}
-			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 accessorAluExp(src$1, ar.#st, unravelAlu(newShape, AluVar.idx));
 			}));
 			const st = ShapeTracker.fromShape(newShape);
-			return new Array$1(exp$3.simplify(), st, exp$3.dtype, backend);
+			arrays.forEach((ar) => ar.dispose());
+			return new Array$1({
+				source: exp$3.simplify(),
+				st,
+				dtype: exp$3.dtype,
+				weakType,
+				backend,
+				committed
+			});
 		}
 		let indices;
 		if (!reduceAxis) indices = unravelAlu(newShape, AluVar.gidx);
@@ -1603,14 +1721,19 @@ var Array$1 = class Array$1 extends Tracer {
 		}
 		const inputs = [];
 		const src = [];
-		for (const ar of arrays) if (ar.#source instanceof AluExp) src.push(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 AluExp) nextSrc = accessorAluExp(ar.#source, ar.#st, indices);
+			else {
+				let gid = inputs.indexOf(ar.#source);
+				if (gid === -1) {
+					gid = inputs.length;
+					inputs.push(ar.#source);
+				}
+				nextSrc = AluExp.globalView(ar.#dtype, gid, ar.#st, indices);
 			}
-			src.push(AluExp.globalView(ar.#dtype, gid, ar.#st, indices));
+			if (!dtypeOverride?.[i]) nextSrc = AluExp.cast(castDtype, nextSrc);
+			src.push(nextSrc);
 		}
 		const exp$2 = custom(src);
 		let re = void 0;
@@ -1623,13 +1746,19 @@ var Array$1 = class Array$1 extends Tracer {
 		const pending = new Set([...arrays.flatMap((ar) => ar.#pending)]);
 		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, ShapeTracker.fromShape(newShape), dtypeOutput, backend, { pending });
+		arrays.forEach((ar) => ar.dispose());
+		return new Array$1({
+			source: output,
+			st: ShapeTracker.fromShape(newShape),
+			dtype: kernel.dtype,
+			weakType,
+			backend,
+			committed,
+			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 Reduction(this.#dtype, op, shape$1[shape$1.length - 1]);
 		const newShape = shape$1.slice(0, -1);
@@ -1648,7 +1777,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, ShapeTracker.fromShape(newShape), this.#dtype, this.#backend, { pending });
+		return this.#newArrayFrom({
+			source: output,
+			st: ShapeTracker.fromShape(newShape),
+			pending
+		});
 	}
 	/**
 	* Normalizes this array into one backed by a `Slot`.
@@ -1684,8 +1817,8 @@ 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, getBackend("cpu"));
+		if (!(this.#source instanceof AluExp)) throw new Error("internal: #dataInline called on non-AluExp source");
+		const ar = this.#newArrayFrom({ backend: getBackend("cpu") });
 		this.dispose();
 		return ar.dataSync();
 	}
@@ -1698,6 +1831,23 @@ var Array$1 = class Array$1 extends Tracer {
 			return ar.#reshape(ar.#st.broadcast(newShape, range(newShape.length - ar.ndim)));
 		});
 	}
+	static #computeBackend(name, arrays) {
+		const committed = arrays.filter((ar) => ar.#committed);
+		if (committed.length > 0) {
+			const backend = committed[0].#backend;
+			for (const ar of committed) if (ar.#backend !== backend) throw new Error(`Device mismatch in ${name} between committed arrays on (${backend.type}, ${ar.#backend.type}), please move to the same device with devicePut()`);
+			return {
+				backend,
+				committed: true
+			};
+		} else {
+			const backend = arrays.length > 0 ? arrays[0].#backend : getBackend();
+			return {
+				backend,
+				committed: false
+			};
+		}
+	}
 	/** Realize the array and return it as data. */
 	async data() {
 		if (this.#source instanceof AluExp && this.size < inlineArrayLimit && this.device !== "cpu") return this.#dataInline();
@@ -1811,7 +1961,11 @@ 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];
 				}
@@ -1853,6 +2007,12 @@ var Array$1 = class Array$1 extends Tracer {
 			[Primitive.Log]([x]) {
 				return [x.#unary(AluOp.Log)];
 			},
+			[Primitive.Erf]([x]) {
+				return [x.#unary(AluOp.Erf)];
+			},
+			[Primitive.Erfc]([x]) {
+				return [x.#unary(AluOp.Erfc)];
+			},
 			[Primitive.Sqrt]([x]) {
 				return [x.#unary(AluOp.Sqrt)];
 			},
@@ -1886,7 +2046,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: 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]) => AluExp.where(cond$1, x$1, y$1);
@@ -1921,7 +2081,8 @@ var Array$1 = class Array$1 extends Tracer {
 			},
 			[Primitive.JitCall](args, { jaxpr, numConsts }) {
 				if (jaxpr.inBinders.length !== args.length) throw new Error(`jit_call expects ${jaxpr.inBinders.length} args, got ${args.length}`);
-				const backend = getBackend();
+				const { backend, committed } = Array$1.#computeBackend("jit_call", args);
+				args = args.map((ar) => ar._putSync(backend));
 				const consts = args.slice(0, numConsts);
 				const tracers = args.slice(numConsts);
 				const jp = jitCompile(backend, jaxpr, consts);
@@ -1932,43 +2093,66 @@ 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, ShapeTracker.fromShape(jaxpr.outs[i].aval.shape), jaxpr.outs[i].aval.dtype, backend, { pending });
+					return new Array$1({
+						source,
+						st: ShapeTracker.fromShape(jaxpr.outs[i].aval.shape),
+						dtype: jaxpr.outs[i].aval.dtype,
+						weakType: jaxpr.outs[i].aval.weakType,
+						backend,
+						committed,
+						pending
+					});
 				});
 			}
 		};
 	}
+	/** @private */
 	_realizeSource() {
 		this.#realize();
 		return this.#source;
 	}
+	/** @private Put this array on a new backend, asynchronously. */
+	async _put(backend) {
+		if (this.#backend === backend) return this;
+		if (this.#source instanceof AluExp) {
+			const ar = this.#newArrayFrom({
+				backend,
+				committed: true
+			});
+			this.dispose();
+			return ar;
+		} else {
+			const data = await this.data();
+			return arrayFromData(data, this.shape, {
+				dtype: this.#dtype,
+				device: backend.type
+			}, this.#weakType);
+		}
+	}
+	/** @private Put this array on a new backend, synchronously. */
+	_putSync(backend) {
+		if (this.#backend === backend) return this;
+		if (this.#source instanceof AluExp) {
+			const ar = this.#newArrayFrom({
+				backend,
+				committed: true
+			});
+			this.dispose();
+			return ar;
+		} else {
+			const data = this.dataSync();
+			return arrayFromData(data, this.shape, {
+				dtype: this.#dtype,
+				device: backend.type
+			}, this.#weakType);
+		}
+	}
 };
-/** Construct an array from a single scalar constant. */
-function scalar(value, { dtype, device } = {}) {
-	if (typeof value === "number") {
-		dtype ??= DType.Float32;
-		if (![
-			DType.Float32,
-			DType.Float16,
-			DType.Int32,
-			DType.Uint32
-		].includes(dtype)) throw new TypeError(`Mismatched dtype for scalar ${value}`);
-	} else if (typeof value === "boolean") {
-		dtype ??= DType.Bool;
-		if (![
-			DType.Float32,
-			DType.Float16,
-			DType.Int32,
-			DType.Uint32,
-			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(AluExp.const(dtype, value), ShapeTracker.fromShape([]), dtype, 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 && !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,
@@ -1990,6 +2174,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 ?? DType.Bool;
 			const data = new Int32Array(flat.map((x) => x ? 1 : 0));
@@ -1998,46 +2186,52 @@ function array(values, { shape: shape$1, dtype, device } = {}) {
 				device
 			});
 		} else {
+			const weakType = dtype == void 0;
 			dtype = dtype ?? DType.Float32;
 			const data = 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 !== DType.Float32) throw new Error("Float32Array must have float32 type");
+		dtype ??= DType.Float32;
+	} else if (data instanceof Int32Array) {
+		if (dtype && dtype !== DType.Int32 && dtype !== DType.Bool) throw new Error("Int32Array must have int32 or bool type");
+		dtype ??= DType.Int32;
+	} else if (data instanceof Uint32Array) {
+		if (dtype && dtype !== DType.Uint32) throw new Error("Uint32Array must have uint32 type");
+		dtype ??= DType.Uint32;
+	} else if (data instanceof Float16Array) {
+		if (dtype && dtype !== DType.Float16) throw new Error("Float16Array must have float16 type");
+		dtype ??= 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 = getBackend(device);
-	if (ArrayBuffer.isView(data)) {
-		const buf = new Uint8Array(data.buffer, data.byteOffset, data.byteLength);
-		if (data instanceof Float32Array) {
-			if (dtype && dtype !== DType.Float32) throw new Error("Float32Array must have float32 type");
-			dtype ??= DType.Float32;
-		} else if (data instanceof Int32Array) {
-			if (dtype && dtype !== DType.Int32 && dtype !== DType.Bool) throw new Error("Int32Array must have int32 or bool type");
-			dtype ??= DType.Int32;
-		} else if (data instanceof Uint32Array) {
-			if (dtype && dtype !== DType.Uint32) throw new Error("Uint32Array must have uint32 type");
-			dtype ??= DType.Uint32;
-		} else if (data instanceof Float16Array) {
-			if (dtype && dtype !== DType.Float16) throw new Error("Float16Array must have float16 type");
-			dtype ??= 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, 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: ShapeTracker.fromShape(shape$1),
+		dtype,
+		weakType,
+		backend,
+		committed: device != void 0
+	});
 }
 function dataToJs(dtype, data, shape$1) {
 	if (shape$1.length === 0) return dtype === DType.Bool ? Boolean(data[0]) : data[0];
@@ -2053,7 +2247,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);
@@ -2064,20 +2258,28 @@ var EvalTrace = class extends Trace {
 };
 const baseArrayTrace = new EvalTrace(newMain(EvalTrace, null));
 const implRules = Array$1._implRules();
+function fullInternal(aval, fillValue, device) {
+	return new Array$1({
+		source: AluExp.const(aval.dtype, fillValue),
+		st: ShapeTracker.fromShape(aval.shape),
+		dtype: aval.dtype,
+		weakType: aval.weakType,
+		backend: getBackend(device),
+		committed: device != void 0
+	});
+}
 function zerosLike$1(val, dtype) {
-	const aval = getAval(val);
-	if (val instanceof Tracer) val.dispose();
-	return zeros(aval.shape, { dtype: dtype ?? aval.dtype });
+	return fullLike(val, 0, dtype);
 }
 function onesLike$1(val, dtype) {
-	const aval = getAval(val);
-	if (val instanceof Tracer) val.dispose();
-	return ones(aval.shape, { dtype: dtype ?? aval.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 } = {}) {
@@ -2095,19 +2297,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 ?? DType.Float32;
-		source = AluExp.const(dtype, fillValue);
-	} else if (typeof fillValue === "bigint") {
-		dtype = dtype ?? DType.Int32;
-		source = AluExp.const(dtype, Number(fillValue));
-	} else if (typeof fillValue === "boolean") {
+	let weakType = dtype == void 0;
+	if (typeof fillValue === "number") dtype = dtype ?? DType.Float32;
+	else if (typeof fillValue === "boolean") {
 		dtype = dtype ?? DType.Bool;
-		source = 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, ShapeTracker.fromShape(shape$1), dtype ?? DType.Float32, getBackend(device));
+	return fullInternal(new ShapedArray(shape$1, dtype, weakType), fillValue, device);
 }
 /**
 * Create an identity matrix.
@@ -2117,6 +2314,7 @@ function full(shape$1, fillValue, { dtype, device } = {}) {
 */
 function eye(numRows, numCols, { dtype, device } = {}) {
 	numCols = numCols ?? numRows;
+	const weakType = dtype == void 0;
 	dtype = dtype ?? DType.Float32;
 	if (numCols < numRows) {
 		const arr = eye(numCols, numRows, {
@@ -2130,7 +2328,14 @@ function eye(numRows, numCols, { dtype, device } = {}) {
 		device
 	});
 	const exp$2 = AluExp.cmplt(AluExp.mod(AluVar.idx, AluExp.i32(numCols + 1)), AluExp.i32(1));
-	return new Array$1(AluExp.cast(dtype, exp$2), ShapeTracker.fromShape([numRows, numCols]), dtype, getBackend(device));
+	return new Array$1({
+		source: AluExp.cast(dtype, exp$2),
+		st: ShapeTracker.fromShape([numRows, numCols]),
+		dtype,
+		weakType,
+		backend: getBackend(device),
+		committed: device != void 0
+	});
 }
 /** Return the identity matrix, with ones on the main diagonal. */
 function identity$1(n, { dtype, device } = {}) {
@@ -2167,7 +2372,14 @@ function arange(start, stop, step = 1, { dtype, device } = {}) {
 	});
 	const exp$2 = AluExp.add(AluExp.const(dtype, start), AluExp.mul(AluExp.cast(dtype, AluVar.idx), AluExp.const(dtype, step)));
 	const st = ShapeTracker.fromShape([size$1]);
-	return new Array$1(exp$2, st, dtype, getBackend(device));
+	return new Array$1({
+		source: exp$2,
+		st,
+		dtype,
+		weakType: false,
+		backend: getBackend(device),
+		committed: device != void 0
+	});
 }
 /**
 * Return evenly spaced numbers over a specified interval.
@@ -2185,10 +2397,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
@@ -2197,7 +2409,14 @@ function linspace(start, stop, num = 50, endpoint = true, { dtype, device } = {}
 	const denom = endpoint ? num - 1 : num;
 	const exp$2 = AluExp.cast(dtype, AluExp.add(AluExp.f32(start), AluExp.mul(AluExp.f32(delta / denom), AluExp.cast(DType.Float32, AluVar.idx))));
 	const st = ShapeTracker.fromShape([num]);
-	return new Array$1(exp$2, st, dtype, getBackend(device));
+	return new Array$1({
+		source: exp$2,
+		st,
+		dtype,
+		weakType: false,
+		backend: getBackend(device),
+		committed: device != void 0
+	});
 }
 function aluCompare(a, b, op) {
 	switch (op) {
@@ -2209,35 +2428,6 @@ function aluCompare(a, b, op) {
 		case CompareOp.LessEqual: return AluExp.add(AluExp.cmplt(a, b), 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/esm/usingCtx.js
@@ -2313,13 +2503,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) {
@@ -2443,14 +2635,19 @@ 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 === 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 === 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;
@@ -2548,7 +2745,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}`);
@@ -2607,7 +2804,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;
 	}
@@ -2676,7 +2873,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]);
@@ -2689,13 +2886,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), DType.Bool)];
+	const aval = promoteAvals(x, y);
+	return [new ShapedArray(aval.shape, DType.Bool, false)];
 }
 function vectorizedUnopAbstractEval([x]) {
 	return [ShapedArray.fromAval(x)];
@@ -2708,18 +2904,18 @@ 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 === DType.Bool || dtype === DType.Bool) throw new TypeError("Bitcast to/from bool is not allowed");
 		if (byteWidth(x.dtype) !== 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 !== DType.Uint32 || k1.dtype !== DType.Uint32) throw new TypeError(`RandomBits requires uint32 keys, got ${k0.dtype} and ${k1.dtype}`);
 		const keyShape = generalBroadcast(k0.shape, k1.shape);
 		if (!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, DType.Uint32)];
+		return [new ShapedArray(shape$1, DType.Uint32, false)];
 	},
 	[Primitive.Sin]: vectorizedUnopAbstractEval,
 	[Primitive.Cos]: vectorizedUnopAbstractEval,
@@ -2727,61 +2923,62 @@ const abstractEvalRules = {
 	[Primitive.Atan]: vectorizedUnopAbstractEval,
 	[Primitive.Exp]: vectorizedUnopAbstractEval,
 	[Primitive.Log]: vectorizedUnopAbstractEval,
+	[Primitive.Erf]: vectorizedUnopAbstractEval,
+	[Primitive.Erfc]: vectorizedUnopAbstractEval,
 	[Primitive.Sqrt]: vectorizedUnopAbstractEval,
 	[Primitive.Min]: binopAbstractEval,
 	[Primitive.Max]: binopAbstractEval,
 	[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 (!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 !== 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 = 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 !== DType.Int32 && a.dtype !== DType.Uint32) throw new TypeError(`Gather indices must be Int32 or Uint32, got ${a.dtype}`);
@@ -2794,7 +2991,7 @@ const abstractEvalRules = {
 		const gatherShape = indices.reduce((shape$1, a) => 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);
@@ -2861,6 +3058,7 @@ function jit$1(f, opts) {
 		const cacheKey = JSON.stringify(jaxprArgs);
 		const { jaxpr, consts, treedef: outTree } = 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
 		});
@@ -2979,6 +3177,16 @@ const jvpRules = {
 	[Primitive.Log]([x], [dx]) {
 		return [[log$1(x.ref)], [reciprocal$1(x).mul(dx)]];
 	},
+	[Primitive.Erf]([x], [dx]) {
+		const coeff = 2 / Math.sqrt(Math.PI);
+		const expTerm = exp$1(neg(x.ref.mul(x.ref)));
+		return [[erf$1(x)], [expTerm.mul(coeff).mul(dx)]];
+	},
+	[Primitive.Erfc]([x], [dx]) {
+		const coeff = -2 / Math.sqrt(Math.PI);
+		const expTerm = exp$1(neg(x.ref.mul(x.ref)));
+		return [[erfc$1(x)], [expTerm.mul(coeff).mul(dx)]];
+	},
 	[Primitive.Sqrt]([x], [dx]) {
 		const z = sqrt$1(x);
 		return [[z.ref], [reciprocal$1(z.mul(2)).mul(dx)]];
@@ -3022,13 +3230,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
 		});
@@ -3082,7 +3291,7 @@ function jvp$1(f, primals, tangents) {
 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) {
@@ -3139,6 +3348,10 @@ var BatchTrace = class extends Trace {
 		const [valsIn, bdimsIn] = unzip2(tracers.map((t) => [t.val, t.batchDim]));
 		const vmapRule = vmapRules[primitive];
 		if (vmapRule === void 0) throw new Error(`No vmap rule for: ${primitive}`);
+		if (bdimsIn.every((d) => d === null)) {
+			const valOuts$1 = bind(primitive, valsIn, params);
+			return valOuts$1.map((x) => new BatchTracer(this, x, null));
+		}
 		const [valOuts, bdimOuts] = vmapRule(this.axisSize, valsIn, bdimsIn, params);
 		return zip(valOuts, bdimOuts).map(([x, bd]) => new BatchTracer(this, x, bd));
 	}
@@ -3146,24 +3359,28 @@ var BatchTrace = class extends Trace {
 		return this.main.globalData;
 	}
 };
-function handleScalarBroadcasting(nd, x, d) {
-	if (d === null || nd === ndim$1(x)) return x;
-	else {
-		const axis = range(ndim$1(x), nd);
-		const shape$1 = [...x.shape, ...axis.map(() => 1)];
-		return broadcast(x, shape$1, axis);
-	}
-}
-/** Process a primitive with built-in broadcasting. */
+/**
+* Process a primitive with built-in broadcasting.
+*
+* Reference: https://github.com/jax-ml/jax/blob/jax-v0.8.1/jax/_src/interpreters/batching.py#L1029
+*/
 function broadcastBatcher(op) {
 	return (axisSize, args, dims) => {
 		if (args.length === 0) throw new Error("Empty list in broadcastBatcher");
-		const idx = dims.findIndex((d) => d !== null);
-		if (idx === -1) return [[op(...args)], [null]];
-		if (zip(args, dims).every(([x, d]) => ndim$1(x) === 0 || deepEqual(x.shape, args[idx].shape) && d === dims[idx])) return [[op(...args)], [dims[idx]]];
-		args = args.map((x, i) => ndim$1(x) > 0 ? moveBatchAxis(axisSize, dims[i], 0, x) : x);
-		const nd = Math.max(...args.map(ndim$1));
-		args = args.map((x, i) => handleScalarBroadcasting(nd, x, dims[i]));
+		const nd = Math.max(...args.map((x, i) => ndim$1(x) + (dims[i] === null ? 1 : 0)));
+		const firstIdx = dims.findIndex((d) => d !== null);
+		const firstBdim = dims[firstIdx] - args[firstIdx].ndim;
+		if (zip(args, dims).every(([x, d]) => d === null && ndim$1(x) < -firstBdim || d !== null && d - x.ndim === firstBdim)) return [[op(...args)], [nd + firstBdim]];
+		args = args.map((x, i) => {
+			if (dims[i] === null) return x;
+			x = moveBatchAxis(axisSize, dims[i], 0, x);
+			if (x.ndim < nd) x = x.reshape([
+				x.shape[0],
+				...rep(nd - x.ndim, 1),
+				...x.shape.slice(1)
+			]);
+			return x;
+		});
 		return [[op(...args)], [0]];
 	};
 }
@@ -3187,17 +3404,18 @@ const vmapRules = {
 	[Primitive.Atan]: unopBatcher(atan$1),
 	[Primitive.Exp]: unopBatcher(exp$1),
 	[Primitive.Log]: unopBatcher(log$1),
+	[Primitive.Erf]: unopBatcher(erf$1),
+	[Primitive.Erfc]: unopBatcher(erfc$1),
 	[Primitive.Sqrt]: unopBatcher(sqrt$1),
 	[Primitive.Min]: broadcastBatcher(min$1),
 	[Primitive.Max]: broadcastBatcher(max$1),
 	[Primitive.Reduce](axisSize, [x], [xBdim], { op, axis }) {
-		if (xBdim === null) return [[reduce(x, op, axis)], [null]];
+		assertNonNull(xBdim);
 		const newAxis = axis.map((ax) => ax + (xBdim <= ax ? 1 : 0));
 		const outBdim = xBdim - axis.filter((ax) => ax < xBdim).length;
 		return [[reduce(x, op, newAxis)], [outBdim]];
 	},
 	[Primitive.Dot](axisSize, [x, y], [xBdim, yBdim]) {
-		if (xBdim === null && yBdim === null) return [[dot$1(x, y)], [null]];
 		x = moveBatchAxis(axisSize, xBdim, x.ndim - (xBdim === null ? 1 : 2), x);
 		y = moveBatchAxis(axisSize, yBdim, y.ndim - (yBdim === null ? 1 : 2), y);
 		const z = dot$1(x, y);
@@ -3206,29 +3424,72 @@ const vmapRules = {
 	[Primitive.Compare](axisSize, args, dims, { op }) {
 		return broadcastBatcher((x, y) => compare(x, y, op))(axisSize, args, dims, {});
 	},
+	[Primitive.Where]: broadcastBatcher(where$1),
+	[Primitive.Transpose](axisSize, [x], [xBdim], { perm }) {
+		assertNonNull(xBdim);
+		const newPerm = perm.map((p) => p + (xBdim <= p ? 1 : 0));
+		newPerm.splice(xBdim, 0, xBdim);
+		return [[transpose$1(x, newPerm)], [xBdim]];
+	},
+	[Primitive.Broadcast](axisSize, [x], [xBdim], { shape: shape$1, axis }) {
+		assertNonNull(xBdim);
+		const newShape = shape$1.toSpliced(xBdim, 0, axisSize);
+		const newAxis = axis.map((ax) => ax + (xBdim <= ax ? 1 : 0));
+		return [[broadcast(x, newShape, newAxis)], [xBdim]];
+	},
 	[Primitive.Reshape](axisSize, [x], [xBdim], { shape: shape$1 }) {
-		if (xBdim === null) return [[reshape$1(x, shape$1)], [null]];
 		x = moveBatchAxis(axisSize, xBdim, 0, x);
 		return [[reshape$1(x, [axisSize, ...shape$1])], [0]];
 	},
 	[Primitive.Flip](axisSize, [x], [xBdim], { axis }) {
-		if (xBdim === null) return [[flip$1(x, axis)], [null]];
+		assertNonNull(xBdim);
 		const newAxis = axis.map((ax) => ax + (xBdim <= ax ? 1 : 0));
 		return [[flip$1(x, newAxis)], [xBdim]];
 	},
 	[Primitive.Shrink](axisSize, [x], [xBdim], { slice }) {
-		if (xBdim === null) return [[shrink(x, slice)], [null]];
+		assertNonNull(xBdim);
 		const newSlice = slice.toSpliced(xBdim, 0, [0, axisSize]);
 		return [[shrink(x, newSlice)], [xBdim]];
 	},
 	[Primitive.Pad](axisSize, [x], [xBdim], { width }) {
-		if (xBdim === null) return [[pad$1(x, width)], [null]];
+		assertNonNull(xBdim);
 		const newWidth = width.toSpliced(xBdim, 0, [0, 0]);
 		return [[pad$1(x, newWidth)], [xBdim]];
 	},
-	[Primitive.JitCall](axisSize, args, dims, { jaxpr }) {
+	[Primitive.Gather](axisSize, [x, ...indices], [xBdim, ...indicesBdim], { axis, outDim }) {
+		if (indicesBdim.every((d) => d === null)) {
+			assertNonNull(xBdim);
+			const newAxis = axis.map((ax) => ax + (xBdim <= ax ? 1 : 0));
+			let newBdim = xBdim - axis.filter((ax) => ax < xBdim).length;
+			let newOutDim = outDim;
+			if (newOutDim < newBdim) newBdim += axis.length;
+			else newOutDim += 1;
+			return [[gather(x, indices, newAxis, newOutDim)], [newBdim]];
+		}
+		const nd = Math.max(...indices.map((m, i) => ndim$1(m) + (indicesBdim[i] === null ? 1 : 0)));
+		indices = indices.map((m, i) => {
+			if (indicesBdim[i] === null) return m;
+			m = moveBatchAxis(axisSize, indicesBdim[i], 0, m);
+			if (m.ndim < nd) m = m.reshape([
+				m.shape[0],
+				...rep(nd - m.ndim, 1),
+				...m.shape.slice(1)
+			]);
+			return m;
+		});
+		if (xBdim === null) return [[gather(x, indices, axis, outDim)], [outDim]];
+		else {
+			x = moveBatchAxis(axisSize, xBdim, 0, x);
+			const newAxis = [0, ...axis.map((ax) => ax + 1)];
+			const extraBatchIndex = arange(axisSize).reshape([-1, ...rep(nd - 1, 1)]);
+			indices.splice(0, 0, extraBatchIndex);
+			return [[gather(x, indices, newAxis, outDim)], [outDim]];
+		}
+	},
+	[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
 		});
@@ -3244,7 +3505,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 = {
@@ -3284,12 +3545,14 @@ function vmapFlat(f, inAxes, args) {
 function vmap$1(f, inAxes = 0) {
 	return (...args) => {
 		const [argsFlat, inTree] = flatten(args);
-		let inAxesFlat;
+		let inAxesFlat = [];
 		if (typeof inAxes === "number") inAxesFlat = rep(argsFlat.length, inAxes);
+		else for (let i = 0; i < args.length; i++) if (inAxes[i] == null) inAxesFlat.push(...rep(inTree.childTreedefs[i].size, null));
+		else if (typeof inAxes[i] === "number") inAxesFlat.push(...rep(inTree.childTreedefs[i].size, inAxes[i]));
 		else {
-			let inTree2;
-			[inAxesFlat, inTree2] = flatten(inAxes);
-			if (!inTree.equals(inTree2)) throw new TreeMismatchError("vmap", inTree, inTree2);
+			const [axesFlat, axesTreeDef] = flatten(inAxes[i]);
+			if (!inTree.childTreedefs[i].equals(axesTreeDef)) throw new TreeMismatchError("vmap", inTree.childTreedefs[i], axesTreeDef);
+			inAxesFlat.push(...axesFlat);
 		}
 		const [fFlat, outTree] = flattenFun(f, inTree);
 		const outsFlat = vmapFlat(fFlat, inAxesFlat, argsFlat);
@@ -3457,8 +3720,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);
@@ -3484,12 +3747,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] = partitionList(inUnknowns, tracers);
 		const outs1Res = bind(Primitive.JitCall, knownTracers.map((t) => t.ref.fullLower()), {
+			name: `${name}_peval`,
 			jaxpr: jaxpr1,
 			numConsts: 0
 		});
@@ -3501,6 +3765,7 @@ var PartialEvalTrace = class extends Trace {
 			prim: Primitive.JitCall,
 			tracersIn: resTracers.concat(unknownTracers),
 			params: {
+				name: `${name}_resid`,
 				jaxpr: jaxpr2,
 				numConsts: 0
 			},
@@ -3643,7 +3908,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}`);
@@ -3823,7 +4088,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]);
@@ -3832,6 +4097,7 @@ const transposeRules = {
 			...residuals,
 			...cts
 		], {
+			name: `${name}_t`,
 			jaxpr: newJaxpr,
 			numConsts: newConsts.length
 		});
@@ -3906,7 +4172,7 @@ 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 (!isFloatDtype(y.dtype)) throw new TypeError("grad only supports floating-point dtypes");
-		const [ct, ...rest] = fVjp(scalar(1, { dtype: y.dtype }));
+		const [ct, ...rest] = fVjp(onesLike$1(y.ref));
 		for (const r of rest) dispose(r);
 		fVjp.dispose();
 		return [y, ct];
@@ -3934,7 +4200,10 @@ __export(lax_exports, {
 	conv: () => conv$1,
 	convGeneralDilated: () => convGeneralDilated,
 	convWithGeneralPadding: () => convWithGeneralPadding,
-	reduceWindow: () => reduceWindow
+	erf: () => erf,
+	erfc: () => erfc,
+	reduceWindow: () => reduceWindow,
+	stopGradient: () => stopGradient$1
 });
 function padtypeToPads(inShape, filterShape, strides, dilation, padding) {
 	const padType = padding.toUpperCase();
@@ -3993,6 +4262,28 @@ function reduceWindow(operand, computation, windowDimensions, windowStrides) {
 		strides: windowStrides
 	}));
 }
+/** The error function: `erf(x) = 2/sqrt(pi) * int[0..x] exp(-t^2) dt`. */
+function erf(x) {
+	return erf$1(x);
+}
+/**
+* The complementary error function: `erfc(x) = 1 - erf(x)`.
+*
+* This function is more accurate than `1 - erf(x)` for large values of `x`,
+* where `erf(x)` is very close to 1.
+*/
+function erfc(x) {
+	return erfc$1(x);
+}
+/**
+* Stops gradient computation.
+*
+* Behaves as the identity function but prevents the flow of gradients during
+* forward or reverse-mode automatic differentiation.
+*/
+function stopGradient$1(x) {
+	return stopGradient(x);
+}
 
 //#endregion
 //#region src/numpy.ts
@@ -4055,6 +4346,9 @@ __export(numpy_exports, {
 	fullLike: () => fullLike$1,
 	greater: () => greater,
 	greaterEqual: () => greaterEqual,
+	hamming: () => hamming,
+	hann: () => hann,
+	heaviside: () => heaviside,
 	hstack: () => hstack,
 	hypot: () => hypot,
 	identity: () => identity$1,
@@ -4276,7 +4570,7 @@ function argmin(a, axis, opts) {
 	} else axis = 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 length = array(shape$1[axis], {
 		dtype: int32,
 		device: a.device
 	});
@@ -4300,7 +4594,7 @@ function argmax(a, axis, opts) {
 	} else axis = 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 length = array(shape$1[axis], {
 		dtype: int32,
 		device: a.device
 	});
@@ -4694,6 +4988,32 @@ function sign(x) {
 	x = fudgeArray(x);
 	return where(notEqual(x.ref, 0), where(less(x.ref, 0), -1, 1), 0);
 }
+/**
+* Return the Hamming window of size M, a taper with a weighted cosine bell.
+*
+* `w(n) = 0.54 - 0.46 * cos(2πn/(M-1))` for `0 <= n <= M-1`.
+*/
+function hamming(M) {
+	return cos(linspace(0, 2 * Math.PI, M)).mul(-.46).add(.54);
+}
+/**
+* Return the Hann window of size M, a taper with a weighted cosine bell.
+*
+* `w(n) = 0.5 - 0.5 * cos(2πn/(M-1))` for `0 <= n <= M-1`.
+*/
+function hann(M) {
+	return cos(linspace(0, 2 * Math.PI, M)).mul(-.5).add(.5);
+}
+/**
+* @function
+* Compute the Heaviside step function. It is defined piecewise:
+* - `heaviside(x1, x2) = 0` for `x1 < 0`,
+* - `heaviside(x1, x2) = x2` for `x1 == 0`,
+* - `heaviside(x1, x2) = 1` for `x1 > 0`.
+*/
+const heaviside = jit$1(function heaviside$1(x1, x2) {
+	return where(less(x1.ref, 0), 0, where(equal(x1, 0), x2, 1));
+});
 /** Calculate element-wise square of the input array. */
 function square(x) {
 	x = fudgeArray(x);
@@ -4713,10 +5033,10 @@ function acos(x) {
 * 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.
+* stable than `sqrt(x1**2 + x2**2)`. We don't currently implement those
+* stability improvements.
 */
-const hypot = jit$1((x1, x2) => {
+const hypot = jit$1(function hypot$1(x1, x2) {
 	return sqrt(square(x1).add(square(x2)));
 });
 /**
@@ -4732,7 +5052,7 @@ const hypot = jit$1((x1, x2) => {
 *
 * The output is ill-defined when both x and y are zero.
 */
-const atan2 = jit$1((y, x) => {
+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);
@@ -4800,13 +5120,13 @@ const degrees = rad2deg;
 * @function
 * Computes first array raised to power of second array, element-wise.
 */
-const power = jit$1((x1, x2) => {
+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((x) => {
+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)));
 });
@@ -4816,7 +5136,7 @@ const cbrt = jit$1((x) => {
 *
 * `sinh(x) = (exp(x) - exp(-x)) / 2`
 */
-const sinh = jit$1((x) => {
+const sinh = jit$1(function sinh$1(x) {
 	const ex = exp(x);
 	const emx = reciprocal(ex.ref);
 	return ex.sub(emx).mul(.5);
@@ -4827,7 +5147,7 @@ const sinh = jit$1((x) => {
 *
 * `cosh(x) = (exp(x) + exp(-x)) / 2`
 */
-const cosh = jit$1((x) => {
+const cosh = jit$1(function cosh$1(x) {
 	const ex = exp(x);
 	const emx = reciprocal(ex.ref);
 	return ex.add(emx).mul(.5);
@@ -4838,7 +5158,7 @@ const cosh = jit$1((x) => {
 *
 * `tanh(x) = sinh(x)/cosh(x) = (exp(x) - exp(-x)) / (exp(x) + exp(-x))`
 */
-const tanh = jit$1((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);
@@ -4849,7 +5169,7 @@ const tanh = jit$1((x) => {
 *
 * `arcsinh(x) = ln(x + sqrt(x^2 + 1))`
 */
-const arcsinh = jit$1((x) => {
+const arcsinh = jit$1(function arcsinh$1(x) {
 	return log(x.ref.add(sqrt(square(x).add(1))));
 });
 /**
@@ -4858,7 +5178,7 @@ const arcsinh = jit$1((x) => {
 *
 * `arccosh(x) = ln(x + sqrt(x^2 - 1))`
 */
-const arccosh = jit$1((x) => {
+const arccosh = jit$1(function arccosh$1(x) {
 	return log(x.ref.add(sqrt(square(x).sub(1))));
 });
 /**
@@ -4867,7 +5187,7 @@ const arccosh = jit$1((x) => {
 *
 * `arctanh(x) = 0.5 * ln((1 + x) / (1 - x))`
 */
-const arctanh = jit$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()`. */
@@ -4983,7 +5303,9 @@ 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
@@ -5036,18 +5358,20 @@ function celu(x, alpha = 1) {
 * @function
 * Gaussion error linear unit (GELU) activation function.
 *
-* This is computed element-wise. Currently jax-js does not support the erf() or
-* gelu() functions exactly as primitives, so an approximation is used:
-* `gelu(x) ~= x * 0.5 * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))`.
+* This is computed element-wise. There are two variants depending on whether
+* `approximate` is set (default true):
 *
-* Reference: https://ml-explore.github.io/mlx/build/html/python/nn/_autosummary_functions/mlx.nn.gelu_approx.html
+* - Approximate: `gelu(x) ~= x * 0.5 * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))`
+* - Exact: `gelu(x) = x * 0.5 * erfc(-x / sqrt(2))`
 *
-* This will be improved in the future.
+* Reference: https://ml-explore.github.io/mlx/build/html/python/nn/_autosummary_functions/mlx.nn.gelu_approx.html
 */
-const gelu = jit$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));
-});
+const gelu = jit$1(function gelu$1(x, opts) {
+	if (opts?.approximate ?? true) {
+		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));
+	} else return x.ref.mul(.5).mul(erfc$1(negative(x.ref.mul(Math.SQRT1_2))));
+}, { staticArgnums: [1] });
 /**
 * Gated linear unit (GLU) activation function.
 *
@@ -5215,8 +5539,11 @@ 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(array(512, {
 		dtype: DType.Uint32,
@@ -5229,7 +5556,7 @@ function uniform(key$1, shape$1 = [], { minval = 0, maxval = 1 } = {}) {
 	const rand = bitcast(float12, 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).
 *
@@ -5240,26 +5567,49 @@ function bernoulli(key$1, p = .5, shape$1 = []) {
 	p = fudgeArray(p);
 	return uniform(key$1, shape$1).less(p);
 }
-/** Sample exponential random values according to `p(x) = exp(-x)`. */
-function exponential(key$1, shape$1 = []) {
+/**
+* @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.
 */
-function normal(key$1, shape$1 = []) {
+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/scipy-special.ts
+var scipy_special_exports = {};
+__export(scipy_special_exports, {
+	erf: () => erf,
+	erfc: () => erfc,
+	logSoftmax: () => logSoftmax,
+	logit: () => logit,
+	logsumexp: () => logsumexp,
+	softmax: () => softmax
+});
+/**
+* @function
+* The logit function, `logit(p) = log(p / (1-p))`.
+*/
+const logit = jit$1(function logit$1(x) {
+	return log(x.ref.div(subtract(1, x)));
+});
 
 //#endregion
 //#region src/polyfills.ts
@@ -5354,6 +5704,25 @@ async function blockUntilReady(x) {
 	await Promise.all(promises);
 	return x;
 }
+/**
+* Transfer `x` to `device`.
+*
+* `x` may be a nested container of arrays or scalars. The resulting structure
+* is committed to the device.
+*
+* If `device` is not specified, this function behaves as identity if the input
+* is already an `Array`, otherwise it places the scalar uncommitted on the
+* default device.
+*/
+async function devicePut(x, device) {
+	const [xflat, structure$1] = flatten(x);
+	const yflat = await Promise.all(xflat.map((leaf) => {
+		if (leaf instanceof Array$1) return device ? leaf._put(getBackend(device)) : Promise.resolve(leaf);
+		else return Promise.resolve(array(leaf, { device }));
+	}));
+	return unflatten(structure$1, yflat);
+}
 
 //#endregion
-export { Array$1 as Array, DType, Jaxpr, blockUntilReady, defaultDevice, devices, grad, init, jacfwd, jacobian, jacrev, jit, jvp, lax_exports as lax, linearize, makeJaxpr, nn_exports as nn, numpy_exports as numpy, random_exports as random, setDebug, tree_exports as tree, valueAndGrad, vjp, vmap };
+export { Array$1 as Array, DType, Jaxpr, blockUntilReady, defaultDevice, devicePut, devices, grad, init, jacfwd, jacobian, jacrev, jit, jvp, lax_exports as lax, linearize, makeJaxpr, nn_exports as nn, numpy_exports as numpy, random_exports as random, scipy_special_exports as scipySpecial, setDebug, tree_exports as tree, valueAndGrad, vjp, vmap };
+//# sourceMappingURL=index.js.map