@jax-js/jax 0.0.2 → 0.0.3

package/dist/index.js

@@ -1,11 +1,12 @@
 import { __export } from "./chunk-Cl8Af3a2.js";
-import { AluExp, AluGroup, AluOp, AluVar, DEBUG, DType, FpHash, Kernel, PPrint, Reduction, ShapeTracker, accessorAluExp, accessorGlobal, byteWidth, checkAxis, deepEqual, devices, dtypedArray, getBackend, init, invertPermutation, isFloatDtype, isNumberPair, isPermutation, partitionList, prod, range, recursiveFlatten, rep, runWithCache, setDevice, toposort, unravelAlu, unzip2, zip } from "./backend-1eVbAoaV.js";
+import { AluExp, AluGroup, AluOp, AluVar, DEBUG, DType, FpHash, Kernel, PPrint, Reduction, ShapeTracker, accessorAluExp, accessorGlobal, byteWidth, checkAxis, deepEqual, devices, dtypedArray, dtypedJsArray, getBackend, init, invertPermutation, isFloatDtype, isNumberPair, isPermutation, partitionList, prod, range, recursiveFlatten, rep, runWithCache, setDevice, toposort, unravelAlu, unzip2, zip, zipn } from "./backend-BqDtPGaR.js";
 
 //#region src/tree.ts
 var tree_exports = {};
 __export(tree_exports, {
 	JsTreeDef: () => JsTreeDef,
 	NodeType: () => NodeType,
+	dispose: () => dispose,
 	flatten: () => flatten,
 	leaves: () => leaves,
 	map: () => map,
@@ -20,7 +21,7 @@ let NodeType = /* @__PURE__ */ function(NodeType$1) {
 	NodeType$1["Leaf"] = "Leaf";
 	return NodeType$1;
 }({});
-/** Analog to the JAX "pytree" object, but for JavaScript. */
+/** Represents the structure of a JsTree. */
 var JsTreeDef = class JsTreeDef {
 	static leaf = new JsTreeDef(NodeType.Leaf, null, []);
 	constructor(nodeType, nodeMetadata, childTreedefs) {
@@ -108,6 +109,194 @@ function map(fn, tree, ...rest) {
 function ref(tree) {
 	return map((x) => x.ref, tree);
 }
+/** Dispose every array in a tree. */
+function dispose(tree) {
+	if (tree) map((x) => x.dispose(), tree);
+}
+
+//#endregion
+//#region src/frontend/convolution.ts
+/**
+* Check that the shapes and parameters passed to convolution are valid.
+*
+* If the check succeeds, returns the output shape.
+*/
+function checkConvShape(lhsShape, rhsShape, { strides, padding, lhsDilation, rhsDilation }) {
+	if (lhsShape.length !== rhsShape.length) throw new Error(`conv() requires inputs with the same number of dimensions, got ${lhsShape.length} and ${rhsShape.length}`);
+	const n = lhsShape.length - 2;
+	if (n < 0) throw new Error("conv() requires at least 2D inputs");
+	if (strides.length !== n) throw new Error("conv() strides != spatial dims");
+	if (padding.length !== n) throw new Error("conv() padding != spatial dims");
+	if (lhsDilation.length !== n) throw new Error("conv() lhsDilation != spatial dimensions");
+	if (rhsDilation.length !== n) throw new Error("conv() rhsDilation != spatial dimensions");
+	if (lhsShape[1] !== rhsShape[1]) throw new Error(`conv() input channels: ${lhsShape[1]} != ${rhsShape[1]}`);
+	const outShape = [lhsShape[0], rhsShape[0]];
+	for (let i = 0; i < n; i++) {
+		if (strides[i] <= 0 || !Number.isInteger(strides[i])) throw new Error(`conv() strides[${i}] must be a positive integer`);
+		if (padding[i].length !== 2 || !padding[i].every(Number.isInteger)) throw new Error(`conv() padding[${i}] must be a 2-tuple of integers`);
+		if (lhsDilation[i] <= 0 || !Number.isInteger(lhsDilation[i])) throw new Error(`conv() lhsDilation[${i}] must be a positive integer`);
+		if (rhsDilation[i] <= 0 || !Number.isInteger(rhsDilation[i])) throw new Error(`conv() rhsDilation[${i}] must be a positive integer`);
+		const [x, k] = [lhsShape[i + 2], rhsShape[i + 2]];
+		if (k <= 0) throw new Error("conv() kernel size must be positive");
+		const [pl, pr] = padding[i];
+		if (pl < -x || pr < -x || pl + pr < -x) throw new Error(`conv() padding[${i}]=(${pl},${pr}) is too negative for input size ${x}`);
+		const kernelSize = (k - 1) * rhsDilation[i] + 1;
+		const inSize = Math.max((x - 1) * lhsDilation[i] + 1, 0) + pl + pr;
+		if (kernelSize > inSize) throw new Error(`conv() kernel size ${kernelSize} > input size ${inSize} in dimension ${i}`);
+		outShape.push(Math.ceil((inSize - kernelSize + 1) / strides[i]));
+	}
+	return outShape;
+}
+function checkPoolShape(inShape, window, strides) {
+	if (strides.length !== window.length) throw new Error("pool() strides != window dims");
+	if (window.length > inShape.length) throw new Error("pool() window has more dimensions than input");
+	const outShape = inShape.slice(0, inShape.length - window.length);
+	for (let i = 0; i < window.length; i++) {
+		const k = window[i];
+		const s = strides[i];
+		const size$1 = inShape[inShape.length - window.length + i];
+		if (k <= 0 || !Number.isInteger(k)) throw new Error(`pool() window[${i}] must be a positive integer`);
+		if (k > size$1) throw new Error(`pool() window[${i}]=${k} > input size ${size$1}`);
+		if (s <= 0 || !Number.isInteger(s)) throw new Error(`pool() strides[${i}] must be a positive integer`);
+		outShape.push(Math.ceil((size$1 - k + 1) / s));
+	}
+	return outShape.concat(window);
+}
+/**
+* Takes a shape tracker and a kernel size `ks`, then reshapes it so the last
+* `ks.length` dimensions become `2 * ks.length` dimensions by treating them as
+* spatial dimensions convolved with a kernel.
+*
+* The resulting array can be multiplied with a kernel of shape `ks`, then
+* reduced along the last `ks.length` dimensions for a convolution.
+*
+* Reference: https://github.com/tinygrad/tinygrad/blob/v0.10.3/tinygrad/tensor.py#L2097
+*/
+function pool(st, ks, strides = 1, dilation = 1) {
+	if (ks.length === 0) return st;
+	if (st.shape.length < ks.length) throw new Error("pool() called with too many dimensions");
+	if (typeof strides === "number") strides = rep(ks.length, strides);
+	if (typeof dilation === "number") dilation = rep(ks.length, dilation);
+	if (strides.some((s) => s <= 0 || !Number.isInteger(s))) throw new Error("pool() strides must be positive integers");
+	if (dilation.some((d) => d <= 0 || !Number.isInteger(d))) throw new Error("pool() dilation must be positive integers");
+	const noop = st.shape.slice(0, -ks.length);
+	const i_ = st.shape.slice(-ks.length);
+	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));
+	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))]);
+	st = st.shrink([...noop.map((x) => [0, x]), ...kidf.map(([k, i, d, f]) => [0, k * (i * f + d)])]).reshape([...noop, ...kidf.flatMap(([k, i, d, f]) => [k, i * f + d])]);
+	const kos = zipn(ks, o_, s_);
+	st = st.shrink([...noop.map((x) => [0, x]), ...kos.flatMap(([k, o, s]) => [[0, k], [0, o * s]])]).reshape([...noop, ...kos.flat(1)]);
+	st = st.shrink([...noop.map((x) => [0, x]), ...kos.flatMap(([k, o]) => [
+		[0, k],
+		[0, o],
+		[0, 1]
+	])]).reshape([...noop, ...kos.flatMap(([k, o]) => [k, o])]);
+	st = st.permute([
+		...range(noop.length),
+		...ks.map((_, j) => noop.length + 2 * j + 1),
+		...ks.map((_, j) => noop.length + 2 * j)
+	]);
+	return st;
+}
+/**
+* Perform the transpose of pool, directly undo-ing a pool() operation.
+*
+* Note that since pool repeats the input, the transpose operation technically
+* should include a sum reduction. This function doesn't perform the reduction,
+* which should be done on the last `k` axes of the returned shape.
+*/
+function poolTranspose(st, inShape, ks, strides = 1, dilation = 1) {
+	if (ks.length === 0) return st;
+	if (typeof strides === "number") strides = rep(ks.length, strides);
+	if (typeof dilation === "number") dilation = rep(ks.length, dilation);
+	const noop = inShape.slice(0, -ks.length);
+	const i_ = inShape.slice(-ks.length);
+	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 (!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_);
+	const kos = zipn(ks, o_, s_);
+	st = st.permute([...range(noop.length), ...ks.flatMap((_, j) => [noop.length + ks.length + j, noop.length + j])]);
+	st = st.reshape([...noop, ...kos.flatMap(([k, o]) => [
+		k,
+		o,
+		1
+	])]).pad([...noop.map(() => [0, 0]), ...s_.flatMap((s) => [
+		[0, 0],
+		[0, 0],
+		[0, s - 1]
+	])]);
+	st = st.reshape([...noop, ...kos.flatMap(([k, o, s]) => [k, o * s])]).pad([...noop.map(() => [0, 0]), ...kidf.flatMap(([_k, i, d, f], j) => [[0, 0], [0, i * f + d - o_[j] * s_[j]]])]);
+	st = st.reshape([...noop, ...kidf.map(([k, i, d, f]) => k * (i * f + d))]).pad([...noop.map(() => [0, 0]), ...kidf.map(([k, i, d, f]) => [0, Math.ceil(k * (i * f + d) / i) * i - k * (i * f + d)])]);
+	st = st.reshape([...noop, ...kidf.flatMap(([k, i, d, f]) => [Math.ceil(k * (i * f + d) / i), i])]).permute([
+		...range(noop.length),
+		...ks.map((_, j) => noop.length + 2 * j + 1),
+		...ks.map((_, j) => noop.length + 2 * j)
+	]);
+	return st;
+}
+/** Applies dilation to an array directly, for transposed convolution. */
+function applyDilation(st, dilation) {
+	if (dilation.every((s) => s === 1)) return st;
+	const s_ = dilation;
+	const [a, b, ...k_] = st.shape;
+	st = st.reshape([
+		a,
+		b,
+		...k_.flatMap((k) => [k, 1])
+	]);
+	st = st.pad([
+		[0, 0],
+		[0, 0],
+		...s_.flatMap((s) => [[0, 0], [0, s - 1]])
+	]);
+	st = st.reshape([
+		a,
+		b,
+		...k_.map((k, i) => k * s_[i])
+	]);
+	st = st.shrink([
+		[0, a],
+		[0, b],
+		...k_.map((k, i) => [0, (k - 1) * s_[i] + 1])
+	]);
+	return st;
+}
+/**
+* Prepare for a convolution between two arrays.
+*
+* This does not check the validity of the shapes, which should be checked
+* beforehand using `checkConvShape()`.
+*/
+function prepareConv(stX, stY, params) {
+	const n = stX.shape.length - 2;
+	stX = applyDilation(stX, params.lhsDilation);
+	const ks = stY.shape.slice(2);
+	stX = stX.padOrShrink([
+		[0, 0],
+		[0, 0],
+		...params.padding
+	]);
+	stX = pool(stX, ks, params.strides, params.rhsDilation);
+	stX = stX.moveaxis(1, n + 1).reshape([
+		stX.shape[0],
+		1,
+		...stX.shape.slice(2, n + 2),
+		stX.shape[1] * prod(ks)
+	]);
+	stY = stY.reshape([
+		stY.shape[0],
+		...rep(n, 1),
+		stY.shape[1] * prod(ks)
+	]);
+	return [stX, stY];
+}
 
 //#endregion
 //#region src/frontend/core.ts
@@ -136,10 +325,14 @@ let Primitive = /* @__PURE__ */ function(Primitive$1) {
 	Primitive$1["Cos"] = "cos";
 	Primitive$1["Exp"] = "exp";
 	Primitive$1["Log"] = "log";
+	Primitive$1["Sqrt"] = "sqrt";
 	Primitive$1["Min"] = "min";
 	Primitive$1["Max"] = "max";
 	Primitive$1["Reduce"] = "reduce";
 	Primitive$1["Dot"] = "dot";
+	Primitive$1["Conv"] = "conv";
+	Primitive$1["Pool"] = "pool";
+	Primitive$1["PoolTranspose"] = "pool_transpose";
 	Primitive$1["Compare"] = "compare";
 	Primitive$1["Where"] = "where";
 	Primitive$1["Transpose"] = "transpose";
@@ -203,6 +396,9 @@ function exp$1(x) {
 function log$1(x) {
 	return bind1(Primitive.Log, [x]);
 }
+function sqrt$1(x) {
+	return bind1(Primitive.Sqrt, [x]);
+}
 function min$1(x, y) {
 	return bind1(Primitive.Min, [x, y]);
 }
@@ -225,6 +421,17 @@ function reduce(x, op, axis, opts) {
 function dot$1(x, y) {
 	return bind1(Primitive.Dot, [x, y]);
 }
+function conv(x, y, params = {}) {
+	if (x.ndim !== y.ndim) throw new Error(`conv() requires inputs with the same number of dimensions, got ${x.ndim} and ${y.ndim}`);
+	const n = x.ndim - 2;
+	if (n < 0) throw new Error("conv() requires at least 2D inputs");
+	return bind1(Primitive.Conv, [x, y], {
+		strides: params.strides ?? rep(n, 1),
+		padding: params.padding ?? rep(n, [0, 0]),
+		lhsDilation: params.lhsDilation ?? rep(n, 1),
+		rhsDilation: params.rhsDilation ?? rep(n, 1)
+	});
+}
 function compare(x, y, op) {
 	return bind1(Primitive.Compare, [x, y], { op });
 }
@@ -360,6 +567,9 @@ var Tracer = class Tracer {
 	get shape() {
 		return this.aval.shape;
 	}
+	get size() {
+		return prod(this.shape);
+	}
 	get dtype() {
 		return this.aval.dtype;
 	}
@@ -411,7 +621,7 @@ var Tracer = class Tracer {
 		else if (typeof axis === "number") axis = [checkAxis(axis, this.ndim)];
 		else axis = axis.map((a) => checkAxis(a, this.ndim));
 		let result = reduce(this, AluOp.Add, axis);
-		result = result.mul(prod(result.shape) / prod(this.shape));
+		result = result.mul(result.size / this.size);
 		if (opts?.keepDims) result = broadcast(result, this.shape, axis);
 		return result;
 	}
@@ -445,8 +655,29 @@ var Tracer = class Tracer {
 	/** Return specified diagonals. See `numpy.diagonal` for full docs. */
 	diagonal(offset = 0, axis1 = 0, axis2 = 1) {
 		if (!Number.isInteger(offset)) throw new TypeError(`offset must be an integer, got ${offset}`);
+		if (offset < 0) return this.diagonal(-offset, axis2, axis1);
+		axis1 = checkAxis(axis1, this.ndim);
+		axis2 = checkAxis(axis2, this.ndim);
 		if (axis1 === axis2) throw new Error("axis1 and axis2 must not be equal");
-		throw new Error("diagonal not implemented");
+		if (offset >= this.shape[axis2]) throw new Error("offset exceeds axis size");
+		let ar = this;
+		if (axis1 !== ar.ndim - 2 || axis2 !== ar.ndim - 1) {
+			const perm = range(ar.ndim).filter((i) => i !== axis1 && i !== axis2).concat(axis1, axis2);
+			ar = ar.transpose(perm);
+		}
+		const [n, m] = ar.shape.slice(-2);
+		const diagSize = Math.min(n, m - offset);
+		ar = ar.reshape([...ar.shape.slice(0, -2), n * m]);
+		const npad = diagSize * (m + 1) - n * m;
+		if (npad > 0) ar = pad$1(ar, [...rep(ar.ndim - 1, [0, 0]), [0, npad]]);
+		else if (npad < 0) ar = shrink(ar, [...ar.shape.slice(0, -1), n * m + npad].map((x) => [0, x]));
+		ar = ar.reshape([
+			...ar.shape.slice(0, -1),
+			diagSize,
+			m + 1
+		]);
+		ar = shrink(ar, [...ar.shape.slice(0, -1).map((x) => [0, x]), [offset, offset + 1]]).reshape(ar.shape.slice(0, -1));
+		return ar;
 	}
 	/** Flatten the array without changing its data. */
 	flatten() {
@@ -589,7 +820,7 @@ var ShapedArray = class ShapedArray {
 	get ndim() {
 		return this.shape.length;
 	}
-	strShort() {
+	toString() {
 		return `${this.dtype}[${this.shape.join(",")}]`;
 	}
 	equals(other) {
@@ -620,7 +851,7 @@ function fullRaise(trace, val) {
 	if (Object.is(val._trace.main, trace.main)) return val;
 	else if (val._trace.main.level < level) return trace.lift(val);
 	else if (val._trace.main.level > level) throw new Error(`Can't lift Tracer level ${val._trace.main.level} to level ${level}`);
-	else throw new Error(`Different traces at same level: ${val._trace}, ${trace}.`);
+	else throw new Error(`Different traces at same level: ${val._trace.constructor}, ${trace.constructor}.`);
 }
 var TreeMismatchError = class extends TypeError {
 	constructor(where$2, left, right) {
@@ -869,16 +1100,16 @@ function jitCompile(backend, jaxpr, consts) {
 	jitCompileCache.set(cacheKey, jp);
 	return jp;
 }
-function reshapeViews(exp$2, mapping) {
+function reshapeViews(exp$2, mapping, reduceAxis = false) {
 	return exp$2.rewrite((exp$3) => {
 		if (exp$3.op === AluOp.GlobalView) {
 			const [gid, st] = exp$3.arg;
 			const newSt = mapping(st);
 			if (newSt) {
-				const indices = unravelAlu(newSt.shape, AluVar.gidx);
+				const indices = reduceAxis ? unravelAlu(newSt.shape.slice(0, -1), AluVar.gidx).concat(AluVar.ridx) : unravelAlu(newSt.shape, AluVar.gidx);
 				return AluExp.globalView(exp$3.dtype, gid, newSt, indices);
 			}
-		}
+		} else if (exp$3.op === AluOp.GlobalIndex) throw new Error("internal: reshapeViews() called with GlobalIndex op");
 	});
 }
 function broadcastedJit(fn) {
@@ -927,6 +1158,7 @@ const jitRules = {
 	[Primitive.Cos]: unopJit(AluExp.cos),
 	[Primitive.Exp]: unopJit(AluExp.exp),
 	[Primitive.Log]: unopJit(AluExp.log),
+	[Primitive.Sqrt]: unopJit(AluExp.sqrt),
 	[Primitive.Min]: broadcastedJit(([a, b]) => AluExp.min(a, b)),
 	[Primitive.Max]: broadcastedJit(([a, b]) => AluExp.max(a, b)),
 	[Primitive.Reduce](nargs, [a], [as], { op, axis }) {
@@ -941,18 +1173,20 @@ const jitRules = {
 		const size$1 = prod(newShape);
 		const reductionSize = prod(shiftedAxes.map((ax) => as.shape[ax]));
 		newShape.push(reductionSize);
-		a = a.rewrite((exp$2) => {
-			if (exp$2.op === AluOp.GlobalView) {
-				const [gid, st] = exp$2.arg;
-				const newSt = st.permute(keptAxes.concat(shiftedAxes)).reshape(newShape);
-				const indices = unravelAlu(newShape.slice(0, -1), AluVar.gidx);
-				indices.push(AluVar.ridx);
-				return AluExp.globalView(exp$2.dtype, gid, newSt, indices);
-			}
-		});
+		const perm = keptAxes.concat(shiftedAxes);
+		a = reshapeViews(a, (st) => st.permute(perm).reshape(newShape), true);
 		const reduction = new Reduction(a.dtype, op, reductionSize);
 		return new Kernel(nargs, size$1, a, reduction);
 	},
+	[Primitive.Pool]: reshapeJit((st, { window, strides }) => pool(st, window, strides)),
+	[Primitive.PoolTranspose](nargs, [a], [as], { inShape, window, strides }) {
+		let stX = poolTranspose(ShapeTracker.fromShape(as.shape), inShape, window, strides);
+		const size$1 = prod(inShape);
+		stX = stX.reshape([...inShape, prod(stX.shape.slice(inShape.length))]);
+		a = reshapeViews(a, (st) => st.compose(stX), true);
+		const reduction = new Reduction(a.dtype, AluOp.Add, stX.shape[stX.shape.length - 1]);
+		return new Kernel(nargs, size$1, a, reduction);
+	},
 	[Primitive.Dot](nargs, [a, b], [as, bs]) {
 		const k1 = jitRules[Primitive.Mul](nargs, [a, b], [as, bs], {});
 		const c = k1.exp;
@@ -962,6 +1196,14 @@ const jitRules = {
 			axis: [cs.ndim - 1]
 		});
 	},
+	[Primitive.Conv](nargs, [a, b], [as, bs], params) {
+		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);
+		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.Transpose]: reshapeJit((st, { perm }) => st.permute(perm)),
@@ -974,8 +1216,20 @@ const jitRules = {
 	}),
 	[Primitive.Shrink]: reshapeJit((st, { slice }) => st.shrink(slice)),
 	[Primitive.Pad]: reshapeJit((st, { width }) => st.pad(width)),
-	[Primitive.Gather]() {
-		throw new Error("Gather is not implemented in JIT yet");
+	[Primitive.Gather](nargs, [x, ...indices], [xs, ...indicesShapes], { axis, outDim }) {
+		const axisSet = new Set(axis);
+		const indexShape = indicesShapes.map((c) => c.shape).reduce(generalBroadcast);
+		const finalShape = xs.shape.filter((_, i) => !axisSet.has(i));
+		finalShape.splice(outDim, 0, ...indexShape);
+		const idxAll = unravelAlu(finalShape, AluVar.gidx);
+		const idxNonaxis = [...idxAll];
+		idxNonaxis.splice(outDim, indexShape.length);
+		const src = [...idxNonaxis];
+		for (let i = 0; i < xs.shape.length; i++) if (axisSet.has(i)) src.splice(i, 0, null);
+		for (const [i, iexp] of indices.entries()) src[axis[i]] = AluExp.cast(DType.Int32, reshapeViews(iexp, (st) => st.broadcast(finalShape, [...range(outDim + indexShape.length - st.shape.length), ...range(outDim + indexShape.length, finalShape.length)])));
+		const [index, valid] = ShapeTracker.fromShape(xs.shape).toAluExp(src);
+		if (!valid.resolve()) throw new Error("internal: expected full validity mask in Gather");
+		return new Kernel(nargs, prod(finalShape), x.substitute({ gidx: index }));
 	},
 	[Primitive.JitCall]() {
 		throw new Error("internal: JitCall should have been flattened before JIT compilation");
@@ -994,9 +1248,15 @@ function splitGraphDataflow(backend, jaxpr) {
 		blackNodes.add(v);
 		p1NextBlack.set(v, v);
 	}
+	const reducePrimitives = [
+		Primitive.Reduce,
+		Primitive.Dot,
+		Primitive.Conv,
+		Primitive.PoolTranspose
+	];
 	for (let i = jaxpr.eqns.length - 1; i >= 0; i--) {
 		const eqn = jaxpr.eqns[i];
-		if (eqn.primitive === Primitive.Reduce || eqn.outBinders.some((v) => blackNodes.has(v))) {
+		if (reducePrimitives.includes(eqn.primitive) || eqn.primitive === Primitive.Gather || eqn.outBinders.some((v) => blackNodes.has(v))) {
 			for (const v of eqn.outBinders) {
 				blackNodes.add(v);
 				p1NextBlack.set(v, v);
@@ -1223,7 +1483,7 @@ var Array$1 = class Array$1 extends Tracer {
 		const inputs = [];
 		const src = [...idxNonaxis];
 		for (let i = 0; i < this.shape.length; i++) if (axisSet.has(i)) src.splice(i, 0, null);
-		for (const [i, ar] of indices.entries()) if (ar.#source instanceof AluExp) src[axis[i]] = AluExp.cast(DType.Int32, accessorAluExp(ar.#dtype, ar.#source, ar.#st, idxAxis));
+		for (const [i, ar] of indices.entries()) if (ar.#source instanceof AluExp) src[axis[i]] = AluExp.cast(DType.Int32, accessorAluExp(ar.#source, ar.#st, idxAxis));
 		else {
 			let gid = inputs.indexOf(ar.#source);
 			if (gid === -1) {
@@ -1233,7 +1493,7 @@ var Array$1 = class Array$1 extends Tracer {
 			src[axis[i]] = AluExp.cast(DType.Int32, AluExp.globalView(ar.#dtype, gid, ar.#st, idxAxis));
 		}
 		let exp$2;
-		if (this.#source instanceof AluExp) exp$2 = accessorAluExp(this.#dtype, this.#source, this.#st, src);
+		if (this.#source instanceof AluExp) exp$2 = accessorAluExp(this.#source, this.#st, src);
 		else {
 			let gid = inputs.indexOf(this.#source);
 			if (gid === -1) {
@@ -1276,7 +1536,7 @@ var Array$1 = class Array$1 extends Tracer {
 		this.#check();
 		if (this.#source instanceof AluExp) {
 			const exp$3 = new AluExp(op, dtypeOutput, [this.#source]);
-			return new Array$1(exp$3, this.#st, dtypeOutput, this.#backend);
+			return new Array$1(exp$3.simplify(), this.#st, dtypeOutput, this.#backend);
 		}
 		const indices = unravelAlu(this.#st.shape, AluVar.gidx);
 		const exp$2 = new AluExp(op, dtypeOutput, [AluExp.globalView(this.#dtype, 0, this.#st, indices)]);
@@ -1309,18 +1569,18 @@ var Array$1 = class Array$1 extends Tracer {
 		if (!dtypeOutput) throw new TypeError("nary operation with no dtype");
 		arrays = Array$1.#broadcastArrays(arrays);
 		const newShape = [...arrays[0].shape];
-		if (arrays.every((ar) => ar.#source instanceof AluExp) && reduceAxis === void 0) {
+		if (arrays.every((ar) => ar.#source instanceof AluExp) && !reduceAxis) {
 			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, arrays[0].#st, exp$4.dtype, backend);
+				return new Array$1(exp$4.simplify(), arrays[0].#st, exp$4.dtype, backend);
 			}
 			const exp$3 = custom(arrays.map((ar) => {
 				const src$1 = ar.#source;
 				if (ar.#st.contiguous) return src$1;
-				return accessorAluExp(ar.#dtype, src$1, ar.#st, unravelAlu(newShape, AluVar.idx));
+				return accessorAluExp(src$1, ar.#st, unravelAlu(newShape, AluVar.idx));
 			}));
 			const st = ShapeTracker.fromShape(newShape);
-			return new Array$1(exp$3, st, exp$3.dtype, backend);
+			return new Array$1(exp$3.simplify(), st, exp$3.dtype, backend);
 		}
 		let indices;
 		if (!reduceAxis) indices = unravelAlu(newShape, AluVar.gidx);
@@ -1330,7 +1590,7 @@ 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.#dtype, ar.#source, ar.#st, indices));
+		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) {
@@ -1364,7 +1624,7 @@ var Array$1 = class Array$1 extends Tracer {
 		const indices = [...unravelAlu(newShape, AluVar.gidx), AluVar.ridx];
 		let exp$2;
 		const inputs = [];
-		if (this.#source instanceof AluExp) exp$2 = accessorAluExp(this.#dtype, this.#source, this.#st, indices);
+		if (this.#source instanceof AluExp) exp$2 = accessorAluExp(this.#source, this.#st, indices);
 		else {
 			inputs.push(this.#source);
 			exp$2 = accessorGlobal(this.#dtype, 0, this.#st, indices);
@@ -1389,7 +1649,7 @@ var Array$1 = class Array$1 extends Tracer {
 		this.#check();
 		const indices = unravelAlu(this.#st.shape, AluVar.gidx);
 		if (this.#source instanceof AluExp) {
-			const exp$2 = accessorAluExp(this.#dtype, this.#source, this.#st, indices);
+			const exp$2 = accessorAluExp(this.#source, this.#st, indices);
 			const kernel = new Kernel(0, this.#st.size, exp$2);
 			const output = this.#backend.malloc(kernel.bytes);
 			const pendingItem = new PendingExecute(this.#backend, kernel, [], [output]);
@@ -1427,42 +1687,51 @@ var Array$1 = class Array$1 extends Tracer {
 	}
 	/** Realize the array and return it as data. */
 	async data() {
-		if (this.#source instanceof AluExp && prod(this.shape) < inlineArrayLimit && this.device !== "cpu") return this.#dataInline();
+		if (this.#source instanceof AluExp && this.size < inlineArrayLimit && this.device !== "cpu") return this.#dataInline();
 		this.#realize();
 		const pending = this.#pending;
 		if (pending) {
 			await Promise.all(pending.map((p) => p.prepare()));
 			for (const p of pending) p.submit();
 		}
-		const byteCount = byteWidth(this.#dtype) * prod(this.shape);
+		const byteCount = byteWidth(this.#dtype) * this.size;
 		const buf = await this.#backend.read(this.#source, 0, byteCount);
 		this.dispose();
 		return dtypedArray(this.dtype, buf);
 	}
-	/** Wait for this array to be placed on the backend, if needed. */
+	/**
+	* Wait for this array to finish evaluation.
+	*
+	* Operations and data loading in jax-js are lazy, so this function ensures
+	* that pending operations are dispatched and fully executed before it
+	* returns.
+	*
+	* If you are mapping from `data()` or `dataSync()`, it will also trigger
+	* dispatch of operations as well.
+	*/
 	async wait() {
 		this.#check();
-		if (this.#source instanceof AluExp) return;
+		if (this.#source instanceof AluExp) return this;
 		const pending = this.#pending;
 		if (pending) {
 			await Promise.all(pending.map((p) => p.prepare()));
 			for (const p of pending) p.submit();
 		}
 		await this.#backend.read(this.#source, 0, 0);
-		this.dispose();
+		return this;
 	}
 	/**
 	* Realize the array and return it as data. This is a sync variant and not
 	* recommended for performance reasons, as it will block rendering.
 	*/
 	dataSync() {
-		if (this.#source instanceof AluExp && prod(this.shape) < inlineArrayLimit && this.device !== "cpu") return this.#dataInline();
+		if (this.#source instanceof AluExp && this.size < inlineArrayLimit && this.device !== "cpu") return this.#dataInline();
 		this.#realize();
 		for (const p of this.#pending) {
 			p.prepareSync();
 			p.submit();
 		}
-		const byteCount = byteWidth(this.#dtype) * prod(this.shape);
+		const byteCount = byteWidth(this.#dtype) * this.size;
 		const buf = this.#backend.readSync(this.#source, 0, byteCount);
 		this.dispose();
 		return dtypedArray(this.dtype, buf);
@@ -1483,6 +1752,16 @@ var Array$1 = class Array$1 extends Tracer {
 	async jsAsync() {
 		return dataToJs(this.dtype, await this.data(), this.shape);
 	}
+	/**
+	* Copy an element of an array to a numeric scalar and return it.
+	*
+	* Throws an error if the array does not have a single element. The array must
+	* either be rank-0, or all dimensions of the shape are 1.
+	*/
+	item() {
+		if (this.size !== 1) throw new Error(`item() can only be called on arrays of size 1`);
+		return this.dataSync()[0];
+	}
 	/** @private Internal plumbing method for Array / Tracer ops. */
 	static _implRules() {
 		return {
@@ -1496,7 +1775,7 @@ var Array$1 = class Array$1 extends Tracer {
 				return [x.#binary(AluOp.Idiv, y)];
 			},
 			[Primitive.Neg]([x]) {
-				return [zerosLike(x).#binary(AluOp.Sub, x)];
+				return [zerosLike(x.ref).#binary(AluOp.Sub, x)];
 			},
 			[Primitive.Reciprocal]([x]) {
 				return [x.#unary(AluOp.Reciprocal)];
@@ -1552,6 +1831,9 @@ var Array$1 = class Array$1 extends Tracer {
 			[Primitive.Log]([x]) {
 				return [x.#unary(AluOp.Log)];
 			},
+			[Primitive.Sqrt]([x]) {
+				return [x.#unary(AluOp.Sqrt)];
+			},
 			[Primitive.Min]([x, y]) {
 				return [x.#binary(AluOp.Min, y)];
 			},
@@ -1562,9 +1844,24 @@ var Array$1 = class Array$1 extends Tracer {
 				if (axis.length === 0) return [x];
 				return [x.#moveAxesDown(axis).#reduce(op)];
 			},
+			[Primitive.Pool]([x], { window, strides }) {
+				const st = pool(x.#st, window, strides);
+				return [x.#reshape(st)];
+			},
+			[Primitive.PoolTranspose]([x], { inShape, window, strides }) {
+				const n = inShape.length;
+				let st = poolTranspose(x.#st, inShape, window, strides);
+				st = st.reshape([...st.shape.slice(0, n), prod(st.shape.slice(n))]);
+				return [x.#reshape(st).#reduce(AluOp.Add)];
+			},
 			[Primitive.Dot]([x, y]) {
 				return [Array$1.#naryCustom("dot", ([x$1, y$1]) => AluExp.mul(x$1, y$1), [x, y], { reduceAxis: true })];
 			},
+			[Primitive.Conv]([x, y], params) {
+				checkConvShape(x.shape, y.shape, params);
+				const [stX, stY] = prepareConv(x.#st, y.#st, params);
+				return [Array$1.#naryCustom("conv", ([x$1, y$1]) => AluExp.mul(x$1, y$1), [x.#reshape(stX), y.#reshape(stY)], { reduceAxis: true })];
+			},
 			[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 })];
@@ -1629,6 +1926,7 @@ function scalar(value, { dtype, device } = {}) {
 		dtype ??= DType.Float32;
 		if (![
 			DType.Float32,
+			DType.Float16,
 			DType.Int32,
 			DType.Uint32
 		].includes(dtype)) throw new TypeError(`Mismatched dtype for scalar ${value}`);
@@ -1636,6 +1934,7 @@ function scalar(value, { dtype, device } = {}) {
 		dtype ??= DType.Bool;
 		if (![
 			DType.Float32,
+			DType.Float16,
 			DType.Int32,
 			DType.Uint32,
 			DType.Bool
@@ -1649,7 +1948,7 @@ function array(values, { shape: shape$1, dtype, device } = {}) {
 		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");
 		return values;
-	} else if (values instanceof Float32Array || values instanceof Int32Array) return arrayFromData(values, shape$1 ?? [values.length], {
+	} else if (ArrayBuffer.isView(values)) return arrayFromData(values, shape$1 ?? [values.length], {
 		dtype,
 		device
 	});
@@ -1678,7 +1977,7 @@ function array(values, { shape: shape$1, dtype, device } = {}) {
 			});
 		} else {
 			dtype = dtype ?? DType.Float32;
-			const data = dtypedArray(dtype, flat);
+			const data = dtypedJsArray(dtype, flat);
 			return arrayFromData(data, shape$1, {
 				dtype,
 				device
@@ -1699,19 +1998,24 @@ function arrayFromData(data, shape$1, { dtype, device } = {}) {
 		});
 	}
 	const backend = getBackend(device);
-	if (data instanceof Float32Array) {
-		if (dtype && dtype !== DType.Float32) throw new Error("Float32Array must have float32 type");
-		const slot = backend.malloc(data.byteLength, data.buffer);
-		return new Array$1(slot, ShapeTracker.fromShape(shape$1), DType.Float32, backend);
-	} else if (data instanceof Int32Array) {
-		if (dtype && dtype !== DType.Int32 && dtype !== DType.Bool) throw new Error("Int32Array must have int32 or bool type");
-		const slot = backend.malloc(data.byteLength, data.buffer);
-		return new Array$1(slot, ShapeTracker.fromShape(shape$1), dtype ?? DType.Int32, backend);
-	} else if (data instanceof Uint32Array) {
-		if (dtype && dtype !== DType.Uint32) throw new Error("Uint32Array must have uint32 type");
-		const slot = backend.malloc(data.byteLength, data.buffer);
-		return new Array$1(slot, ShapeTracker.fromShape(shape$1), DType.Uint32, backend);
-	} else throw new Error("Unsupported data type");
+	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);
 }
 function dataToJs(dtype, data, shape$1) {
 	if (shape$1.length === 0) return dtype === DType.Bool ? Boolean(data[0]) : data[0];
@@ -1738,9 +2042,20 @@ var EvalTrace = class extends Trace {
 };
 const baseArrayTrace = new EvalTrace(newMain(EvalTrace, null));
 const implRules = Array$1._implRules();
-function zerosLike(val) {
+function zerosLike(val, dtype) {
+	const aval = getAval(val);
+	if (val instanceof Tracer) val.dispose();
+	return zeros(aval.shape, { dtype: dtype ?? aval.dtype });
+}
+function onesLike(val, dtype) {
+	const aval = getAval(val);
+	if (val instanceof Tracer) val.dispose();
+	return ones(aval.shape, { dtype: dtype ?? aval.dtype });
+}
+function fullLike(val, fillValue, dtype) {
 	const aval = getAval(val);
-	return zeros(aval.shape, { dtype: aval.dtype });
+	if (val instanceof Tracer) val.dispose();
+	return full(aval.shape, fillValue, { dtype: dtype ?? aval.dtype });
 }
 /** Return a new array of given shape and type, filled with zeros. */
 function zeros(shape$1, { dtype, device } = {}) {
@@ -1762,6 +2077,9 @@ function full(shape$1, fillValue, { dtype, device } = {}) {
 	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") {
 		dtype = dtype ?? DType.Bool;
 		source = AluExp.const(dtype, fillValue ? 1 : 0);
@@ -1859,7 +2177,6 @@ function linspace(start, stop, num = 50, endpoint = true, { dtype, device } = {}
 	const st = ShapeTracker.fromShape([num]);
 	return new Array$1(exp$2, st, dtype, getBackend(device));
 }
-/** Translate a `CompareOp` into an `AluExp` on two sub-expressions. */
 function aluCompare(a, b, op) {
 	switch (op) {
 		case CompareOp.Greater: return AluExp.mul(AluExp.cmpne(a, b), AluExp.cmplt(a, b).not());
@@ -1901,7 +2218,7 @@ function generalBroadcast(a, b) {
 }
 
 //#endregion
-//#region node_modules/.pnpm/@oxc-project+runtime@0.77.3/node_modules/@oxc-project/runtime/src/helpers/esm/usingCtx.js
+//#region node_modules/.pnpm/@oxc-project+runtime@0.78.0/node_modules/@oxc-project/runtime/src/helpers/esm/usingCtx.js
 function _usingCtx() {
 	var r = "function" == typeof SuppressedError ? SuppressedError : function(r$1, e$2) {
 		var n$1 = Error();
@@ -1969,7 +2286,7 @@ var Var = class Var {
 		this.aval = aval;
 	}
 	toString() {
-		return `Var(${this.id}):${this.aval.strShort()}`;
+		return `Var(${this.id}):${this.aval.toString()}`;
 	}
 };
 /** Literal in a Jaxpr expression. Currently, only scalars are supported. */
@@ -2009,7 +2326,7 @@ var VarPrinter = class {
 		return name;
 	}
 	nameType(v) {
-		return `${this.name(v)}:${v.aval.strShort()}`;
+		return `${this.name(v)}:${v.aval.toString()}`;
 	}
 };
 /** A single statement / binding in a Jaxpr, in ANF form. */
@@ -2164,8 +2481,8 @@ var JaxprType = class {
 		this.outTypes = outTypes;
 	}
 	toString() {
-		const inTypes = this.inTypes.map((aval) => aval.strShort()).join(", ");
-		const outTypes = this.outTypes.map((aval) => aval.strShort()).join(", ");
+		const inTypes = this.inTypes.map((aval) => aval.toString()).join(", ");
+		const outTypes = this.outTypes.map((aval) => aval.toString()).join(", ");
 		return `(${inTypes}) -> (${outTypes})`;
 	}
 };
@@ -2244,7 +2561,7 @@ var JaxprTracer = class extends Tracer {
 		this.aval = aval;
 	}
 	toString() {
-		return `JaxprTracer(${this.aval.strShort()})`;
+		return `JaxprTracer(${this.aval.toString()})`;
 	}
 	get ref() {
 		return this;
@@ -2383,6 +2700,7 @@ const abstractEvalRules = {
 	[Primitive.Cos]: vectorizedUnopAbstractEval,
 	[Primitive.Exp]: vectorizedUnopAbstractEval,
 	[Primitive.Log]: vectorizedUnopAbstractEval,
+	[Primitive.Sqrt]: vectorizedUnopAbstractEval,
 	[Primitive.Min]: binopAbstractEval,
 	[Primitive.Max]: binopAbstractEval,
 	[Primitive.Reduce]([x], { axis }) {
@@ -2390,6 +2708,15 @@ const abstractEvalRules = {
 		const newShape = x.shape.filter((_, i) => !axisSet.has(i));
 		return [new ShapedArray(newShape, x.dtype)];
 	},
+	[Primitive.Pool]([x], { window, strides }) {
+		const shape$1 = checkPoolShape(x.shape, window, strides);
+		return [new ShapedArray(shape$1, x.dtype)];
+	},
+	[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)];
+	},
 	[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");
@@ -2397,6 +2724,11 @@ const abstractEvalRules = {
 		shape$1.splice(-1, 1);
 		return [new ShapedArray(shape$1, x.dtype)];
 	},
+	[Primitive.Conv]([lhs, rhs], params) {
+		if (lhs.dtype !== rhs.dtype) throw new TypeError(`Conv dtype mismatch, got ${lhs.dtype} vs ${rhs.dtype}`);
+		const shape$1 = checkConvShape(lhs.shape, rhs.shape, params);
+		return [new ShapedArray(shape$1, lhs.dtype)];
+	},
 	[Primitive.Compare]: compareAbstractEval,
 	[Primitive.Where]([cond, x, y]) {
 		if (cond.dtype !== DType.Bool) throw new TypeError(`Condition must be boolean, got ${cond.dtype}`);
@@ -2444,15 +2776,34 @@ const abstractEvalRules = {
 		return outTypes;
 	}
 };
-function makeJaxpr$1(f) {
+function splitIdx(values, argnums) {
+	const a = [];
+	const b = [];
+	for (let i = 0; i < values.length; i++) if (argnums.has(i)) a.push(values[i]);
+	else b.push(values[i]);
+	return [a, b];
+}
+function joinIdx(n, a, b, argnums) {
+	const result = [];
+	let ai = 0;
+	let bi = 0;
+	for (let i = 0; i < n; i++) if (argnums.has(i)) result.push(a[ai++]);
+	else result.push(b[bi++]);
+	return result;
+}
+function makeJaxpr$1(f, opts) {
 	return (...argsIn) => {
 		try {
 			var _usingCtx$1 = _usingCtx();
-			const [avalsIn, inTree] = flatten(argsIn);
-			const [fFlat, outTree] = flattenFun(f, inTree);
+			const staticArgnums = new Set(opts?.staticArgnums ?? []);
+			const [staticArgs, shapedArgs] = splitIdx(argsIn, staticArgnums);
+			const [avalsIn, inTree] = flatten(shapedArgs);
+			const [fFlat, outTree] = flattenFun((...dynamicArgs) => {
+				return f(...joinIdx(argsIn.length, staticArgs, dynamicArgs, staticArgnums));
+			}, inTree);
 			const builder = new JaxprBuilder();
 			const main = _usingCtx$1.u(newMain(JaxprTrace, builder));
-			const _dynamic = _usingCtx$1.u(newDynamic(main));
+			_usingCtx$1.u(newDynamic(main));
 			const trace = new JaxprTrace(main);
 			const tracersIn = avalsIn.map((aval) => trace.newArg(typeof aval === "object" ? aval : pureArray(aval)));
 			const outs = fFlat(...tracersIn);
@@ -2471,14 +2822,17 @@ function makeJaxpr$1(f) {
 		}
 	};
 }
-function jit$1(f) {
+function jit$1(f, opts) {
 	const cache = /* @__PURE__ */ new Map();
+	const staticArgnums = new Set(opts?.staticArgnums ?? []);
 	return ((...args) => {
-		const [argsFlat, inTree] = flatten(args);
+		const [staticArgs, dynamicArgs] = splitIdx(args, staticArgnums);
+		const [argsFlat, inTree] = flatten(dynamicArgs);
 		const avalsInFlat = argsFlat.map((x) => ShapedArray.fromAval(getAval(x)));
 		const avalsIn = unflatten(inTree, avalsInFlat);
-		const cacheKey = JSON.stringify(avalsIn);
-		const { jaxpr, consts, treedef: outTree } = runWithCache(cache, cacheKey, () => makeJaxpr$1(f)(...avalsIn));
+		const jaxprArgs = joinIdx(args.length, staticArgs, avalsIn, staticArgnums);
+		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], {
 			jaxpr,
 			numConsts: consts.length
@@ -2515,7 +2869,7 @@ var JVPTrace = class extends Trace {
 		return this.lift(pureArray(val));
 	}
 	lift(val) {
-		return new JVPTracer(this, val, zerosLike(val));
+		return new JVPTracer(this, val, zerosLike(val.ref));
 	}
 	processPrimitive(primitive, tracers, params) {
 		const [primalsIn, tangentsIn] = unzip2(tracers.map((x) => [x.primal, x.tangent]));
@@ -2533,19 +2887,25 @@ function linearTangentsJvp(primitive) {
 		return [ys, dys];
 	};
 }
+/** Rule for product of gradients in bilinear operations. */
+function bilinearTangentsJvp(primitive) {
+	return ([x, y], [dx, dy], params) => {
+		const primal = bind1(primitive, [x.ref, y.ref], params);
+		const tangent = bind1(primitive, [x, dy], params).add(bind1(primitive, [dx, y], params));
+		return [[primal], [tangent]];
+	};
+}
 /** Rule that zeros out any tangents. */
 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))];
+		return [ys, ys.map((y) => zerosLike(y.ref))];
 	};
 }
 const jvpRules = {
 	[Primitive.Add]: linearTangentsJvp(Primitive.Add),
-	[Primitive.Mul]([x, y], [dx, dy]) {
-		return [[x.ref.mul(y.ref)], [x.mul(dy).add(dx.mul(y))]];
-	},
+	[Primitive.Mul]: bilinearTangentsJvp(Primitive.Mul),
 	[Primitive.Idiv]: zeroTangentsJvp(Primitive.Idiv),
 	[Primitive.Neg]: linearTangentsJvp(Primitive.Neg),
 	[Primitive.Reciprocal]([x], [dx]) {
@@ -2558,13 +2918,13 @@ const jvpRules = {
 		if (isFloatDtype(dtype) && isFloatDtype(x.dtype)) return [[cast(x, dtype)], [cast(dx, dtype)]];
 		else {
 			dx.dispose();
-			return [[cast(x, dtype)], [zerosLike(x)]];
+			return [[cast(x.ref, dtype)], [zerosLike(x)]];
 		}
 	},
 	[Primitive.Bitcast]([x], [dx], { dtype }) {
 		if (x.dtype === dtype) return [[x], [dx]];
 		dx.dispose();
-		return [[bitcast(x, dtype)], [zerosLike(x)]];
+		return [[bitcast(x.ref, dtype)], [zerosLike(x)]];
 	},
 	[Primitive.RandomBits]: zeroTangentsJvp(Primitive.RandomBits),
 	[Primitive.Sin]([x], [dx]) {
@@ -2580,6 +2940,10 @@ const jvpRules = {
 	[Primitive.Log]([x], [dx]) {
 		return [[log$1(x.ref)], [reciprocal$1(x).mul(dx)]];
 	},
+	[Primitive.Sqrt]([x], [dx]) {
+		const z = sqrt$1(x);
+		return [[z.ref], [reciprocal$1(z.mul(2)).mul(dx)]];
+	},
 	[Primitive.Min]([x, y], [dx, dy]) {
 		return [[min$1(x.ref, y.ref)], [where$1(less$1(y, x), dy, dx)]];
 	},
@@ -2596,13 +2960,14 @@ const jvpRules = {
 			const primal = reduce(x.ref, op, axis);
 			const notMin = notEqual$1(x, broadcast(primal.ref, x.shape, axis));
 			const minCount = where$1(notMin.ref, 0, 1).sum(axis);
-			const tangent = where$1(notMin, op === AluOp.Min ? 0 : 0, dx).sum(axis).mul(reciprocal$1(minCount));
+			const tangent = where$1(notMin, 0, dx).sum(axis).div(minCount);
 			return [[primal], [tangent]];
 		} else throw new Error(`JVP rule not implemented for reduce op: ${op}`);
 	},
-	[Primitive.Dot]([x, y], [dx, dy]) {
-		return [[dot$1(x.ref, y.ref)], [dot$1(dx, y).add(dot$1(x, dy))]];
-	},
+	[Primitive.Pool]: linearTangentsJvp(Primitive.Pool),
+	[Primitive.PoolTranspose]: linearTangentsJvp(Primitive.PoolTranspose),
+	[Primitive.Dot]: bilinearTangentsJvp(Primitive.Dot),
+	[Primitive.Conv]: bilinearTangentsJvp(Primitive.Conv),
 	[Primitive.Compare]: zeroTangentsJvp(Primitive.Compare),
 	[Primitive.Where]([cond, x, y], [dcond, dx, dy]) {
 		dcond.dispose();
@@ -2778,6 +3143,7 @@ const vmapRules = {
 	[Primitive.Cos]: unopBatcher(cos$1),
 	[Primitive.Exp]: unopBatcher(exp$1),
 	[Primitive.Log]: unopBatcher(log$1),
+	[Primitive.Sqrt]: unopBatcher(sqrt$1),
 	[Primitive.Min]: broadcastBatcher(min$1),
 	[Primitive.Max]: broadcastBatcher(max$1),
 	[Primitive.Reduce](axisSize, [x], [xBdim], { op, axis }) {
@@ -2914,7 +3280,7 @@ var PartialVal = class PartialVal {
 		return this.val !== null;
 	}
 	toString() {
-		return this.val ? this.val.toString() : this.aval.strShort();
+		return this.val ? this.val.toString() : this.aval.toString();
 	}
 };
 function partialEvalFlat(f, pvalsIn) {
@@ -3288,12 +3654,72 @@ const transposeRules = {
 		if (op === AluOp.Add) return [broadcast(ct, x.aval.shape, axis)];
 		else throw new NonlinearError(Primitive.Reduce);
 	},
+	[Primitive.Pool]([ct], [x], { window, strides }) {
+		if (!(x instanceof UndefPrimal)) throw new NonlinearError(Primitive.Pool);
+		return bind(Primitive.PoolTranspose, [ct], {
+			inShape: x.aval.shape,
+			window,
+			strides
+		});
+	},
+	[Primitive.PoolTranspose]([ct], [x], { window, strides }) {
+		if (!(x instanceof UndefPrimal)) throw new NonlinearError(Primitive.PoolTranspose);
+		return bind(Primitive.Pool, [ct], {
+			window,
+			strides
+		});
+	},
 	[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];
 		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];
 	},
+	[Primitive.Conv]([ct], [lhs, rhs], params) {
+		if (lhs instanceof UndefPrimal === rhs instanceof UndefPrimal) throw new NonlinearError(Primitive.Conv);
+		const rev01 = [
+			1,
+			0,
+			...range(2, ct.ndim)
+		];
+		if (lhs instanceof UndefPrimal) {
+			let kernel = rhs;
+			kernel = transpose$1(kernel, rev01);
+			kernel = flip$1(kernel, range(2, kernel.ndim));
+			const result = conv(ct, kernel, {
+				strides: params.lhsDilation,
+				padding: params.padding.map(([pl, _pr], i) => {
+					const dilatedKernel = (kernel.shape[i + 2] - 1) * params.rhsDilation[i] + 1;
+					const dilatedCt = (ct.shape[i + 2] - 1) * params.strides[i] + 1;
+					const padBefore = dilatedKernel - 1 - pl;
+					const dilatedLhs = (lhs.aval.shape[i + 2] - 1) * params.lhsDilation[i] + 1;
+					const padAfter = dilatedLhs + dilatedKernel - 1 - dilatedCt - padBefore;
+					return [padBefore, padAfter];
+				}),
+				lhsDilation: params.strides,
+				rhsDilation: params.rhsDilation
+			});
+			return [result, null];
+		} else {
+			const newLhs = transpose$1(lhs, rev01);
+			const newRhs = transpose$1(ct, rev01);
+			let result = conv(newLhs, newRhs, {
+				strides: params.rhsDilation,
+				padding: params.padding.map(([pl, _pr], i) => {
+					const dilatedLhs = (lhs.aval.shape[i + 2] - 1) * params.lhsDilation[i] + 1;
+					const dilatedKernel = (rhs.aval.shape[i + 2] - 1) * params.rhsDilation[i] + 1;
+					const dilatedCt = (ct.shape[i + 2] - 1) * params.strides[i] + 1;
+					const padFromLhs = dilatedCt - dilatedLhs;
+					const padFromRhs = dilatedKernel - pl - 1;
+					return [pl, padFromLhs + padFromRhs];
+				}),
+				lhsDilation: params.lhsDilation,
+				rhsDilation: params.strides
+			});
+			result = transpose$1(result, rev01);
+			return [null, result];
+		}
+	},
 	[Primitive.Where]([ct], [cond, x, y]) {
 		const cts = [
 			null,
@@ -3414,8 +3840,8 @@ function valueAndGrad$1(f) {
 		if (x.length === 0) throw new Error("grad requires at least one argument to differentiate");
 		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 (y.dtype !== DType.Float32) throw new TypeError("grad currently only supports float32");
-		const [ct, ...rest] = fVjp(pureArray(1));
+		if (!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();
 		return [y, ct];
 	};
@@ -3429,6 +3855,73 @@ function jacrev$1(f) {
 	};
 }
 
+//#endregion
+//#region src/lax.ts
+var lax_exports = {};
+__export(lax_exports, {
+	conv: () => conv$1,
+	convGeneralDilated: () => convGeneralDilated,
+	convWithGeneralPadding: () => convWithGeneralPadding,
+	reduceWindow: () => reduceWindow
+});
+function padtypeToPads(inShape, filterShape, strides, dilation, padding) {
+	const padType = padding.toUpperCase();
+	switch (padType) {
+		case "VALID": return rep(inShape.length, [0, 0]);
+		case "SAME":
+		case "SAME_LOWER": {
+			const outShape = inShape.map((size$1, i) => Math.ceil(size$1 / strides[i]));
+			const padSizes = zipn(outShape, strides, filterShape, dilation, inShape).map(([o, s, k, d, i]) => Math.max(0, (o - 1) * s + 1 + (k - 1) * d - i));
+			if (padType === "SAME") return padSizes.map((size$1) => [size$1 >> 1, size$1 - (size$1 >> 1)]);
+			else return padSizes.map((size$1) => [size$1 - (size$1 >> 1), size$1 >> 1]);
+		}
+		default: throw new Error(`Unknown padding type: ${padType}`);
+	}
+}
+/**
+* General n-dimensional convolution operator, with optional dilation.
+*
+* The semantics of this operation mimic the `jax.lax.conv_general_dilated`
+* function in JAX, which wraps XLA's general convolution operator.
+*
+* Grouped convolutions are not supported right now.
+*/
+function convGeneralDilated(lhs, rhs, windowStrides, padding, { lhsDilation, rhsDilation } = {}) {
+	if (lhs.ndim < 2) throw new Error("lhs must have at least 2 dimensions");
+	if (rhs.ndim < 2) throw new Error("rhs must have at least 2 dimensions");
+	if (typeof padding === "string") {
+		if (lhsDilation?.some((d) => d !== 1)) throw new Error("String padding is not supported for transposed convolutions");
+		padding = padtypeToPads(lhs.shape.slice(2), rhs.shape.slice(2), windowStrides, rhsDilation ?? rep(rhs.ndim - 2, 1), padding);
+	}
+	return conv(lhs, rhs, {
+		strides: windowStrides,
+		padding,
+		lhsDilation,
+		rhsDilation
+	});
+}
+/** Convenience wrapper around `convGeneralDilated`. */
+function convWithGeneralPadding(lhs, rhs, windowStrides, padding, lhsDilation, rhsDilation) {
+	return convGeneralDilated(lhs, rhs, windowStrides, padding, {
+		lhsDilation,
+		rhsDilation
+	});
+}
+/** Convenience wrapper around `convGeneralDilated`. */
+function conv$1(lhs, rhs, windowStrides, padding) {
+	return convGeneralDilated(lhs, rhs, windowStrides, padding);
+}
+/** Reduce a computation over padded windows. */
+function reduceWindow(operand, computation, windowDimensions, windowStrides) {
+	if (operand.ndim < windowDimensions.length) throw new Error(`Operand dimensions ${operand.ndim} < window ${windowDimensions.length}`);
+	if (!windowStrides) windowStrides = rep(windowDimensions.length, 1);
+	for (let i = 0; i < operand.ndim; i++) computation = vmap$1(computation, 0);
+	return computation(bind1(Primitive.Pool, [operand], {
+		window: windowDimensions,
+		strides: windowStrides
+	}));
+}
+
 //#endregion
 //#region src/numpy.ts
 var numpy_exports = {};
@@ -3447,9 +3940,9 @@ __export(numpy_exports, {
 	bool: () => bool,
 	clip: () => clip,
 	columnStack: () => columnStack,
-	complex64: () => complex64,
 	concatenate: () => concatenate,
 	cos: () => cos,
+	cosh: () => cosh,
 	diag: () => diag,
 	diagonal: () => diagonal,
 	divide: () => divide,
@@ -3464,8 +3957,10 @@ __export(numpy_exports, {
 	flip: () => flip,
 	fliplr: () => fliplr,
 	flipud: () => flipud,
+	float16: () => float16,
 	float32: () => float32,
 	full: () => full,
+	fullLike: () => fullLike$1,
 	greater: () => greater,
 	greaterEqual: () => greaterEqual,
 	hstack: () => hstack,
@@ -3492,6 +3987,7 @@ __export(numpy_exports, {
 	negative: () => negative,
 	notEqual: () => notEqual,
 	ones: () => ones,
+	onesLike: () => onesLike$1,
 	pad: () => pad,
 	permuteDims: () => permuteDims,
 	pi: () => pi,
@@ -3502,11 +3998,14 @@ __export(numpy_exports, {
 	scalar: () => scalar,
 	shape: () => shape,
 	sin: () => sin,
+	sinh: () => sinh,
 	size: () => size,
+	sqrt: () => sqrt,
 	square: () => square,
 	stack: () => stack,
 	sum: () => sum,
 	tan: () => tan,
+	tanh: () => tanh,
 	transpose: () => transpose,
 	trueDivide: () => trueDivide,
 	trunc: () => trunc,
@@ -3515,13 +4014,14 @@ __export(numpy_exports, {
 	vecdot: () => vecdot,
 	vstack: () => vstack,
 	where: () => where,
-	zeros: () => zeros
+	zeros: () => zeros,
+	zerosLike: () => zerosLike$1
 });
 const float32 = DType.Float32;
 const int32 = DType.Int32;
 const uint32 = DType.Uint32;
 const bool = DType.Bool;
-const complex64 = DType.Complex64;
+const float16 = DType.Float16;
 /** Euler's constant, `e = 2.7182818284590...` */
 const e = Math.E;
 /** Euler-Mascheroni constant, `γ = 0.5772156649...` */
@@ -3548,6 +4048,8 @@ const cos = cos$1;
 const exp = exp$1;
 /** 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. */
+const sqrt = sqrt$1;
 /** Return element-wise minimum of the input arrays. */
 const minimum = min$1;
 /** Return element-wise maximum of the input arrays. */
@@ -3589,6 +4091,12 @@ const pad = pad$1;
 const ndim = ndim$1;
 /** 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. */
+const fullLike$1 = fullLike;
 /**
 * Return the number of elements in an array, optionally along an axis.
 * Does not consume array reference.
@@ -3639,13 +4147,7 @@ function argmin(a, axis, opts) {
 		dtype: int32,
 		device: a.device
 	});
-	const idx = where(isMax, scalar(1, {
-		dtype: int32,
-		device: a.device
-	}), scalar(0, {
-		dtype: int32,
-		device: a.device
-	})).mul(arange(shape$1[axis], 0, -1, {
+	const idx = isMax.astype(DType.Int32).mul(arange(shape$1[axis], 0, -1, {
 		dtype: int32,
 		device: a.device
 	}).reshape([shape$1[axis], ...rep(shape$1.length - axis - 1, 1)]));
@@ -3669,13 +4171,7 @@ function argmax(a, axis, opts) {
 		dtype: int32,
 		device: a.device
 	});
-	const idx = where(isMax, scalar(1, {
-		dtype: int32,
-		device: a.device
-	}), scalar(0, {
-		dtype: int32,
-		device: a.device
-	})).mul(arange(shape$1[axis], 0, -1, {
+	const idx = isMax.astype(DType.Int32).mul(arange(shape$1[axis], 0, -1, {
 		dtype: int32,
 		device: a.device
 	}).reshape([shape$1[axis], ...rep(shape$1.length - axis - 1, 1)]));
@@ -3807,9 +4303,11 @@ function ravel(a) {
 * Return specified diagonals.
 *
 * If a is 2D, return the diagonal of the array with the given offset. If a is
-* 3D or higher, compute diagonals along the two given axes.
+* 3D or higher, compute diagonals along the two given axes (default: 0, 1).
 *
-* This returns a view over the existing array.
+* This returns a view over the existing array. The shape of the resulting array
+* is determined by removing the two axes along which the diagonal is taken,
+* then appending a new axis to the right with holding the diagonals.
 */
 function diagonal(a, offset, axis1, axis2) {
 	return fudgeArray(a).diagonal(offset, axis1, axis2);
@@ -3825,15 +4323,16 @@ 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, 0);
-		if (k !== 0) throw new Error("diag() for 1D arrays only for k=0");
-		return ret;
+		const ret = where(eye(n).equal(1), a.ref, zerosLike$1(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;
 	} else if (a.ndim === 2) return diagonal(a, k);
 	else throw new TypeError("numpy.diag only supports 1D and 2D arrays");
 }
 /** Return if two arrays are element-wise equal within a tolerance. */
 function allclose(actual, expected, options) {
-	const { rtol = 1e-5, atol = 1e-8 } = options ?? {};
+	const { rtol = 1e-5, atol = 1e-7 } = options ?? {};
 	const x = array(actual);
 	const y = array(expected);
 	if (!deepEqual(x.shape, y.shape)) return false;
@@ -3967,15 +4466,52 @@ function log2(x) {
 function log10(x) {
 	return log(x).mul(Math.LOG10E);
 }
+/**
+* Calculate element-wise hyperbolic sine of input.
+*
+* `sinh(x) = (exp(x) - exp(-x)) / 2`
+*/
+function sinh(x) {
+	const ex = exp(x);
+	const emx = reciprocal(ex.ref);
+	return ex.sub(emx).mul(.5);
+}
+/**
+* Calculate element-wise hyperbolic cosine of input.
+*
+* `cosh(x) = (exp(x) + exp(-x)) / 2`
+*/
+function cosh(x) {
+	const ex = exp(x);
+	const emx = reciprocal(ex.ref);
+	return ex.add(emx).mul(.5);
+}
+/**
+* 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 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);
+}
 
 //#endregion
 //#region src/nn.ts
 var nn_exports = {};
 __export(nn_exports, {
+	celu: () => celu,
+	elu: () => elu,
+	gelu: () => gelu,
+	glu: () => glu,
 	identity: () => identity,
+	leakyRelu: () => leakyRelu,
 	logSigmoid: () => logSigmoid,
 	logSoftmax: () => logSoftmax,
 	logsumexp: () => logsumexp,
+	mish: () => mish,
 	oneHot: () => oneHot,
 	relu: () => relu,
 	relu6: () => relu6,
@@ -4035,10 +4571,7 @@ function softSign(x) {
 *
 * Reference: https://en.wikipedia.org/wiki/Swish_function
 */
-function silu(x) {
-	x = fudgeArray(x);
-	return x.ref.mul(sigmoid(x));
-}
+const silu = jit$1((x) => x.ref.mul(sigmoid(x)));
 /**
 * Sigmoid-weighted Linear Unit (SiLU) activation function, also known as
 * Swish, computed element-wise:
@@ -4058,6 +4591,72 @@ function logSigmoid(x) {
 }
 /** Identity activation function. Returns the argument unmodified. */
 const identity = fudgeArray;
+/** Leaky rectified linear (ReLU) activation function */
+function leakyRelu(x, negativeSlope = .01) {
+	x = fudgeArray(x);
+	return where(less(x.ref, 0), x.ref.mul(negativeSlope), x);
+}
+/**
+* Exponential linear unit activation function.
+*
+* Computes the element-wise function:
+* `elu(x) = x > 0 ? x : alpha * (exp(x) - 1)`
+*/
+function elu(x, alpha = 1) {
+	x = fudgeArray(x);
+	return where(less(x.ref, 0), exp(x.ref).sub(1).mul(alpha), x);
+}
+/**
+* Continuously-differentiable exponential linear unit activation function.
+*
+* Computes the element-wise function:
+* `celu(x) = x > 0 ? x : alpha * (exp(x/alpha) - 1)`
+*/
+function celu(x, alpha = 1) {
+	x = fudgeArray(x);
+	return where(less(x.ref, 0), exp(x.ref.div(alpha)).sub(1).mul(alpha), x);
+}
+/**
+* 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)))`.
+*
+* Reference: https://ml-explore.github.io/mlx/build/html/python/nn/_autosummary_functions/mlx.nn.gelu_approx.html
+*
+* This will be improved in the future.
+*/
+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));
+});
+/**
+* Gated linear unit (GLU) activation function.
+*
+* Splits the `axis` dimension of the input into two halves, a and b, then
+* computes `a * sigmoid(b)`.
+*/
+function glu(x, axis = -1) {
+	x = fudgeArray(x);
+	axis = checkAxis(axis, x.ndim);
+	const size$1 = x.shape[axis];
+	if (size$1 % 2 !== 0) throw new Error(`glu: axis ${axis} of shape (${x.shape}) does not have even length`);
+	const slice = x.shape.map((a$1) => [0, a$1]);
+	const a = shrink(x.ref, slice.toSpliced(axis, 1, [0, size$1 / 2]));
+	const b = shrink(x, slice.toSpliced(axis, 1, [size$1 / 2, size$1]));
+	return a.mul(sigmoid(b));
+}
+/**
+* Mish activation function.
+*
+* Computes the element-wise function:
+* `mish(x) = x * tanh(softplus(x))`
+*/
+function mish(x) {
+	x = fudgeArray(x);
+	return x.ref.mul(tanh(softplus(x)));
+}
 /**
 * Softmax function. Computes the function which rescales elements to the range
 * [0, 1] such that the elements along `axis` sum to 1.
@@ -4134,7 +4733,7 @@ function logsumexp(x, axis) {
 * ```
 */
 function oneHot(x, numClasses) {
-	if (x.dtype !== "int32") throw new TypeError(`oneHot expects integers, got ${x.dtype}`);
+	if (x.dtype !== DType.Int32) throw new TypeError(`oneHot expects integers, got ${x.dtype}`);
 	return eye(numClasses, void 0, { device: x.device }).slice(x);
 }
 
@@ -4203,6 +4802,19 @@ const vmap = vmap$1;
 const jacfwd = jacfwd$1;
 /** Construct a Jaxpr by dynamically tracing a function with example inputs. */
 const makeJaxpr = makeJaxpr$1;
+/**
+* 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.
+*
+* **Options:**
+* - `staticArgnums`: An array of argument indices to treat as static
+*   (compile-time constant). These arguments must be hashable, won't be traced,
+*   and different values will trigger recompilation.
+* - `device`: The device to place the computation on. If not specified, the
+*   computation will be placed on the default device.
+*/
 const jit = jit$1;
 /**
 * Produce a local linear approximation to a function at a point using jvp() and
@@ -4224,4 +4836,4 @@ const jacrev = jacrev$1;
 const jacobian = jacrev;
 
 //#endregion
-export { devices, grad, init, jacfwd, jacobian, jacrev, jit, jvp, linearize, makeJaxpr, nn_exports as nn, numpy_exports as numpy, random_exports as random, setDevice, tree_exports as tree, valueAndGrad, vjp, vmap };
+export { DType, 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, setDevice, tree_exports as tree, valueAndGrad, vjp, vmap };