@jax-js/jax 0.1.8 → 0.1.10

package/dist/index.cjs

@@ -30,7 +30,7 @@ var __toESM = (mod$1, isNodeMode, target) => (target = mod$1 != null ? __create(
 }) : target, mod$1));
 
 //#endregion
-const require_backend = require('./backend-B3foXiV_.cjs');
+const require_backend = require('./backend-DMauYnfl.cjs');
 
 //#region src/frontend/convolution.ts
 /**
@@ -838,6 +838,11 @@ var Tracer = class Tracer {
 		if (this.dtype === dtype) return this;
 		return cast(this, dtype);
 	}
+	/** Return a bitwise cast of the array, viewed as a new dtype. */
+	view(dtype) {
+		if (!dtype || dtype === this.dtype) return this;
+		return bitcast(this, dtype);
+	}
 	/** Subtract an array from this one. */
 	sub(other) {
 		return this.add(neg(other));
@@ -920,18 +925,25 @@ var Tracer = class Tracer {
 		return sort$1(this.transpose(perm)).transpose(require_backend.invertPermutation(perm));
 	}
 	/**
-	* Return the indices that would sort an array. This may not be a stable
-	* sorting algorithm; it need not preserve order of indices in ties.
+	* Return the indices that would sort an array. Unlike `sort`, this is
+	* guaranteed to be a stable sorting algorithm; it always returns the smaller
+	* index first in event of ties.
 	*
 	* See `jax.numpy.argsort` for full docs.
 	*/
 	argsort(axis = -1) {
 		axis = require_backend.checkAxis(axis, this.ndim);
-		if (axis === this.ndim - 1) return argsort$1(this)[1];
+		if (axis === this.ndim - 1) {
+			const [y$1, yi$1] = argsort$1(this);
+			y$1.dispose();
+			return yi$1;
+		}
 		const perm = require_backend.range(this.ndim);
 		perm.splice(axis, 1);
 		perm.push(axis);
-		return argsort$1(this.transpose(perm))[1].transpose(require_backend.invertPermutation(perm));
+		const [y, yi] = argsort$1(this.transpose(perm));
+		y.dispose();
+		return yi.transpose(require_backend.invertPermutation(perm));
 	}
 	/**
 	* Slice an array along one or more axes.
@@ -1652,7 +1664,7 @@ const abstractEvalRules = {
 		return [new ShapedArray(x.shape, dtype, false)];
 	},
 	[Primitive.Bitcast]([x], { dtype }) {
-		if (x.dtype === require_backend.DType.Bool || dtype === require_backend.DType.Bool) throw new TypeError("Bitcast to/from bool is not allowed");
+		if (x.dtype !== dtype && (x.dtype === require_backend.DType.Bool || dtype === require_backend.DType.Bool)) throw new TypeError("Bitcast to/from bool is not allowed");
 		if (require_backend.byteWidth(x.dtype) !== require_backend.byteWidth(dtype)) throw new TypeError(`Bitcast from ${x.dtype} to ${dtype} with different byte width`);
 		return [new ShapedArray(x.shape, dtype, false)];
 	},
@@ -3074,8 +3086,8 @@ var Array$1 = class Array$1 extends Tracer {
 				return [x.#unary(require_backend.AluOp.Cast, dtype)];
 			},
 			[Primitive.Bitcast]([x], { dtype }) {
-				if (x.dtype === require_backend.DType.Bool || dtype === require_backend.DType.Bool) throw new TypeError("Bitcast to/from bool is not allowed");
 				if (x.dtype === dtype) return [x];
+				if (x.dtype === require_backend.DType.Bool || dtype === require_backend.DType.Bool) throw new TypeError("Bitcast to/from bool is not allowed");
 				if (require_backend.byteWidth(x.dtype) !== require_backend.byteWidth(dtype)) throw new TypeError(`Bitcast from ${x.dtype} to ${dtype} with different byte width`);
 				if (x.#source instanceof require_backend.AluExp) return [x.#unary(require_backend.AluOp.Bitcast, dtype)];
 				else {
@@ -3416,32 +3428,26 @@ function fullInternal(aval, fillValue, device) {
 		committed: device != void 0
 	});
 }
-function zerosLike$1(val, dtype) {
-	return fullLike(val, 0, dtype);
+function zerosLike$1(val, opts) {
+	return fullLike(val, 0, opts);
 }
-function onesLike$1(val, dtype) {
-	return fullLike(val, 1, dtype);
+function onesLike$1(val, opts) {
+	return fullLike(val, 1, opts);
 }
-function fullLike(val, fillValue, dtype) {
+function fullLike(val, fillValue, { dtype, shape: shape$1, device } = {}) {
 	const aval = getAval(val);
 	if (val instanceof Tracer) val.dispose();
 	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);
+	const sa = new ShapedArray(shape$1 ?? aval.shape, dtype ?? aval.dtype, aval.weakType && dtype === void 0);
+	return fullInternal(sa, fillValue, device);
 }
 /** Return a new array of given shape and type, filled with zeros. */
-function zeros(shape$1, { dtype, device } = {}) {
-	return full(shape$1, 0, {
-		dtype,
-		device
-	});
+function zeros(shape$1, opts) {
+	return full(shape$1, 0, opts);
 }
 /** Return a new array of given shape and type, filled with ones. */
-function ones(shape$1, { dtype, device } = {}) {
-	return full(shape$1, 1, {
-		dtype,
-		device
-	});
+function ones(shape$1, opts) {
+	return full(shape$1, 1, opts);
 }
 /** Return a new array of given shape and type, filled with `fill_value`. */
 function full(shape$1, fillValue, { dtype, device } = {}) {
@@ -4178,6 +4184,7 @@ const jvpRules = {
 	},
 	[Primitive.TriangularSolve]([a, b], [da, db], { unitDiagonal }) {
 		const x = triangularSolve$1(a.ref, b, { unitDiagonal });
+		da = unitDiagonal ? triu(da, 1) : triu(da);
 		const dax = batchMatmulT(da, x.ref);
 		const rhsT = db.sub(mT(dax));
 		const dx = triangularSolve$1(a, rhsT, { unitDiagonal });
@@ -5253,6 +5260,7 @@ function ifft(a, axis = -1) {
 var numpy_linalg_exports = {};
 __export(numpy_linalg_exports, {
 	cholesky: () => cholesky,
+	cross: () => cross$1,
 	det: () => det,
 	diagonal: () => diagonal,
 	inv: () => inv,
@@ -5283,6 +5291,19 @@ function cholesky(a, { upper = false, symmetrizeInput = true } = {}) {
 	if (symmetrizeInput) a = a.ref.add(matrixTranspose(a)).mul(.5);
 	return cholesky$1(a, { upper });
 }
+/**
+* Compute the cross-product of two 3D vectors.
+*
+* This is a simpler and less flexible version of `jax.numpy.cross()`.
+* Both inputs must have size 3 along the specified axis.
+*/
+function cross$1(x1, x2, axis = -1) {
+	const a1 = require_backend.checkAxis(axis, ndim(x1));
+	const a2 = require_backend.checkAxis(axis, ndim(x2));
+	if (shape(x1)[a1] !== 3) throw new Error(`linalg.cross: x1 must have size 3 along axis ${axis}, got ${shape(x1)[a1]}`);
+	if (shape(x2)[a2] !== 3) throw new Error(`linalg.cross: x2 must have size 3 along axis ${axis}, got ${shape(x2)[a2]}`);
+	return cross(x1, x2, { axis });
+}
 /** Compute the determinant of a square matrix (batched). */
 function det(a) {
 	a = fudgeArray(a);
@@ -5298,7 +5319,7 @@ function det(a) {
 function inv(a) {
 	a = fudgeArray(a);
 	const n = checkSquare("inv", a);
-	return solve(a, eye(n));
+	return solve(a, eye(n, void 0, { dtype: a.dtype }));
 }
 /**
 * Return the least-squares solution to a linear equation.
@@ -5332,7 +5353,7 @@ function lstsq(a, b) {
 			lower: true,
 			transposeA: true
 		});
-		return matmul(at, llb.ref);
+		return matmul(at, llb);
 	} else {
 		const ata = matmul(at.ref, a);
 		const l = cholesky(ata, { symmetrizeInput: false });
@@ -5355,8 +5376,9 @@ function matrixPower(a, n) {
 	a = fudgeArray(a);
 	const m = checkSquare("matrixPower", a);
 	if (n === 0) {
+		const dtype = a.dtype;
 		a.dispose();
-		return broadcastTo(eye(m), a.shape);
+		return broadcastTo(eye(m, void 0, { dtype }), a.shape);
 	}
 	if (n < 0) {
 		a = inv(a);
@@ -5423,7 +5445,7 @@ function solve(a, b) {
 		lower: true,
 		unitDiagonal: true
 	});
-	let x = triangularSolve(lu$2, LPb.ref, {
+	let x = triangularSolve(lu$2, LPb, {
 		leftSide: true,
 		lower: false
 	});
@@ -5538,13 +5560,17 @@ __export(numpy_exports, {
 	argmax: () => argmax,
 	argmin: () => argmin,
 	argsort: () => argsort,
+	around: () => round,
 	array: () => array,
+	arrayEqual: () => arrayEqual,
+	arrayEquiv: () => arrayEquiv,
 	asin: () => asin,
 	asinh: () => arcsinh,
 	astype: () => astype,
 	atan: () => atan,
 	atan2: () => atan2,
 	atanh: () => arctanh,
+	average: () => average,
 	bool: () => bool,
 	broadcastArrays: () => broadcastArrays,
 	broadcastShapes: () => broadcastShapes,
@@ -5555,11 +5581,13 @@ __export(numpy_exports, {
 	columnStack: () => columnStack,
 	concatenate: () => concatenate,
 	convolve: () => convolve,
+	copysign: () => copysign,
 	corrcoef: () => corrcoef,
 	correlate: () => correlate,
 	cos: () => cos,
 	cosh: () => cosh,
 	cov: () => cov,
+	cross: () => cross,
 	cumsum: () => cumsum,
 	cumulativeSum: () => cumsum,
 	deg2rad: () => deg2rad,
@@ -5595,7 +5623,6 @@ __export(numpy_exports, {
 	fullLike: () => fullLike$1,
 	greater: () => greater,
 	greaterEqual: () => greaterEqual,
-	hamming: () => hamming,
 	hann: () => hann,
 	heaviside: () => heaviside,
 	hstack: () => hstack,
@@ -5619,9 +5646,14 @@ __export(numpy_exports, {
 	log10: () => log10,
 	log1p: () => log1p,
 	log2: () => log2,
+	logicalAnd: () => logicalAnd,
+	logicalNot: () => logicalNot,
+	logicalOr: () => logicalOr,
+	logicalXor: () => logicalXor,
 	logspace: () => logspace,
 	matmul: () => matmul,
 	matrixTranspose: () => matrixTranspose,
+	matvec: () => matvec,
 	max: () => max,
 	maximum: () => maximum,
 	mean: () => mean,
@@ -5654,6 +5686,8 @@ __export(numpy_exports, {
 	remainder: () => remainder,
 	repeat: () => repeat,
 	reshape: () => reshape,
+	rint: () => rint,
+	round: () => round,
 	shape: () => shape,
 	sign: () => sign,
 	sin: () => sin,
@@ -5686,6 +5720,7 @@ __export(numpy_exports, {
 	var_: () => var_,
 	vdot: () => vdot,
 	vecdot: () => vecdot,
+	vecmat: () => vecmat,
 	vstack: () => vstack,
 	where: () => where,
 	zeros: () => zeros,
@@ -5749,6 +5784,22 @@ const notEqual = notEqual$1;
 const greaterEqual = greaterEqual$1;
 /** @function Compare two arrays element-wise. */
 const lessEqual = lessEqual$1;
+/** Compute element-wise logical AND. */
+function logicalAnd(x, y) {
+	return astype(x, require_backend.DType.Bool).mul(astype(y, require_backend.DType.Bool));
+}
+/** Compute element-wise logical OR. */
+function logicalOr(x, y) {
+	return astype(x, require_backend.DType.Bool).add(astype(y, require_backend.DType.Bool));
+}
+/** Compute element-wise logical XOR. */
+function logicalXor(x, y) {
+	return notEqual(astype(x, require_backend.DType.Bool), astype(y, require_backend.DType.Bool));
+}
+/** Compute element-wise logical NOT. */
+function logicalNot(x) {
+	return notEqual(astype(x, require_backend.DType.Bool), true);
+}
 /** @function Element-wise ternary operator, evaluates to `x` if cond else `y`. */
 const where = where$1;
 /**
@@ -5856,6 +5907,34 @@ function mean(a, axis = null, opts) {
 	return fudgeArray(a).mean(axis, opts);
 }
 /**
+* Compute the weighted average along the specified axis.
+*
+* If no axis is specified, mean is computed along all the axes. The weights
+* should have shape matching that of `a`, or if an axis is specified, it should
+* match the shape along those axes.
+*/
+function average(a, axis = null, opts) {
+	a = fudgeArray(a);
+	if (opts?.weights == null) return mean(a, axis, opts);
+	const weights = fudgeArray(opts.weights);
+	axis = require_backend.normalizeAxis(axis, ndim(a));
+	const wShape = weights.shape;
+	const aShape = a.shape;
+	if (require_backend.deepEqual(wShape, aShape)) {
+		const scl = sum(weights.ref, axis, opts);
+		return sum(multiply(a, weights), axis, opts).div(scl);
+	} else if (axis.length === 1 && wShape.length === 1 && wShape[0] === aShape[axis[0]]) {
+		const broadcastShape = aShape.map((_, i) => i === axis[0] ? wShape[0] : 1);
+		const wReshaped = reshape(weights, broadcastShape);
+		const scl = sum(wReshaped.ref, axis, opts);
+		return sum(multiply(a, wReshaped), axis, opts).div(scl);
+	} else {
+		weights.dispose();
+		a.dispose();
+		throw new Error(`average: weights shape ${JSON.stringify(wShape)} is not compatible with array shape ${JSON.stringify(aShape)} and axis ${JSON.stringify(axis)}`);
+	}
+}
+/**
 * Returns the indices of the minimum values along an axis.
 *
 * By default, index is into the flatted array, otherwise it is along the
@@ -6234,8 +6313,9 @@ function sort(a, axis = -1) {
 	return fudgeArray(a).sort(axis);
 }
 /**
-* Return indices that would sort an array. This may be an unstable sorting
-* algorithm; it need not preserve order of indices in ties.
+* Return indices that would sort an array. Unlike `sort`, this is guaranteed to
+* be a stable sorting algorithm; it always returns the smaller index first in
+* event of ties.
 *
 * Returns an array of `int32` indices.
 *
@@ -6258,20 +6338,63 @@ function take(a, indices, axis = null) {
 	axis = require_backend.checkAxis(axis, ndim(a));
 	return gather(a, [indices], [axis], axis);
 }
-/** Return if two arrays are element-wise equal within a tolerance. */
+/**
+* Return if two arrays are element-wise equal within a tolerance.
+*
+* The formula used is `|actual - expected| <= atol + rtol * |expected|`, with
+* NaN values comparing equal if `equalNaN` is true.
+*/
 function allclose(actual, expected, options) {
-	const { rtol = 1e-5, atol = 1e-7 } = options ?? {};
+	const { rtol = 1e-5, atol = 1e-7, equalNaN = false } = options ?? {};
 	const x = array(actual);
 	const y = array(expected);
 	if (!require_backend.deepEqual(x.shape, y.shape)) return false;
 	const xData = x.dataSync();
 	const yData = y.dataSync();
 	for (let i = 0; i < xData.length; i++) {
-		if (isNaN(xData[i]) !== isNaN(yData[i])) return false;
+		if (equalNaN ? isNaN(xData[i]) !== isNaN(yData[i]) : isNaN(xData[i]) || isNaN(yData[i])) return false;
 		if (Math.abs(xData[i] - yData[i]) > atol + rtol * Math.abs(yData[i])) return false;
 	}
 	return true;
 }
+/**
+* Check if two arrays are element-wise equal.
+*
+* Returns False if the arrays have different shapes. If `equalNaN` is True,
+* NaNs in the same position are considered equal.
+*/
+function arrayEqual(a1, a2, opts) {
+	a1 = fudgeArray(a1);
+	a2 = fudgeArray(a2);
+	if (!require_backend.deepEqual(a1.shape, a2.shape)) {
+		a1.dispose();
+		a2.dispose();
+		return array(false);
+	}
+	if (opts?.equalNaN) {
+		const nanMask = isnan(a1.ref).mul(isnan(a2.ref));
+		return where(nanMask, true, equal(a1, a2)).all();
+	}
+	return equal(a1, a2).all();
+}
+/**
+* Check if two arrays are element-wise equal after broadcasting.
+*
+* Unlike `arrayEqual`, this allows inputs with different but
+* broadcast-compatible shapes.
+*/
+function arrayEquiv(a1, a2) {
+	a1 = fudgeArray(a1);
+	a2 = fudgeArray(a2);
+	try {
+		const [b1, b2] = broadcastArrays(a1, a2);
+		return equal(b1, b2).all();
+	} catch {
+		a1.dispose();
+		a2.dispose();
+		return array(false);
+	}
+}
 /** Matrix product of two arrays. */
 function matmul(x, y) {
 	if (ndim(x) === 0 || ndim(y) === 0) throw new Error("matmul: x and y must be at least 1D");
@@ -6285,6 +6408,16 @@ function matmul(x, y) {
 		rhsBatchDims: require_backend.range(-2 - numBatchDims, -2)
 	});
 }
+/** Matrix-vector product. x1 is [..., M, N], x2 is [..., N] → [..., M]. */
+function matvec(x1, x2) {
+	if (ndim(x1) < 2 || ndim(x2) < 1) throw new Error("matvec: x1 must be at least 2D and x2 at least 1D");
+	return einsum("...mn,...n->...m", x1, x2);
+}
+/** Vector-matrix product. x1 is [..., N], x2 is [..., N, M] → [..., M]. */
+function vecmat(x1, x2) {
+	if (ndim(x1) < 1 || ndim(x2) < 2) throw new Error("vecmat: x1 must be at least 1D and x2 at least 2D");
+	return einsum("...n,...nm->...m", x1, x2);
+}
 /** Dot product of two arrays. */
 function dot$1(x, y) {
 	if (ndim(x) === 0 || ndim(y) === 0) return multiply(x, y);
@@ -6443,6 +6576,49 @@ function outer(x, y) {
 	y = ravel(y);
 	return multiply(x.reshape([x.shape[0], 1]), y);
 }
+/**
+* @function Compute the cross product of two arrays.
+*
+* Supports 2D (scalar result) and 3D cross products, with optional axis
+* arguments. If `axis` is given, it overrides `axisa`, `axisb`, and `axisc`.
+*/
+const cross = jit$1(function cross$2(a, b, { axisa = -1, axisb = -1, axisc = -1, axis } = {}) {
+	if (axis !== void 0) {
+		axisa = axis;
+		axisb = axis;
+		axisc = axis;
+	}
+	axisa = require_backend.checkAxis(axisa, ndim(a));
+	axisb = require_backend.checkAxis(axisb, ndim(b));
+	a = moveaxis$1(a, axisa, -1);
+	b = moveaxis$1(b, axisb, -1);
+	const da = a.shape.at(-1);
+	const db = b.shape.at(-1);
+	if (da !== 2 && da !== 3 || db !== 2 && db !== 3) throw new Error(`cross: incompatible dimensions for cross product (got ${da} and ${db})`);
+	if (da === 2 && db === 2) {
+		const [a0$1, a1$1] = split$1(a, 2, -1);
+		const [b0$1, b1$1] = split$1(b, 2, -1);
+		return squeeze(a0$1.mul(b1$1).sub(a1$1.mul(b0$1)), -1);
+	}
+	if (da === 2) {
+		const zeroShape = [...a.shape.slice(0, -1), 1];
+		a = concatenate([a, zeros(zeroShape)], -1);
+	}
+	if (db === 2) {
+		const zeroShape = [...b.shape.slice(0, -1), 1];
+		b = concatenate([b, zeros(zeroShape)], -1);
+	}
+	const [a0, a1, a2] = split$1(a, 3, -1);
+	const [b0, b1, b2] = split$1(b, 3, -1);
+	const c0 = a1.ref.mul(b2.ref).sub(a2.ref.mul(b1.ref));
+	const c1 = a2.mul(b0.ref).sub(a0.ref.mul(b2));
+	const c2 = a0.mul(b1).sub(a1.mul(b0));
+	return moveaxis$1(concatenate([
+		c0,
+		c1,
+		c2
+	], -1), -1, axisc);
+}, { staticArgnums: [2] });
 /** Vector dot product of two arrays along a given axis. */
 function vecdot(x, y, { axis } = {}) {
 	const xaxis = require_backend.checkAxis(axis ?? -1, ndim(x));
@@ -6537,18 +6713,17 @@ function absolute(x) {
 /** Return an element-wise indication of sign of the input. */
 function sign(x) {
 	x = fudgeArray(x);
-	return where(notEqual(x.ref, 0), where(less(x.ref, 0), -1, 1), 0);
+	return where(notEqual(x.ref, 0), where(less(x, 0), -1, 1), 0);
 }
-/** @function Return element-wise positive values of the input (no-op). */
-const positive = fudgeArray;
 /**
-* 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
+* Return the value with the magnitude of x and the sign of y, element-wise.
 */
-function hamming(M) {
-	return cos(linspace(0, 2 * Math.PI, M)).mul(-.46).add(.54);
-}
+const copysign = jit$1(function copysign$1(x, y) {
+	return absolute(x).mul(sign(y));
+});
+/** @function Return element-wise positive values of the input (no-op). */
+const positive = fudgeArray;
 /**
 * Return the Hann window of size M, a taper with a weighted cosine bell.
 *
@@ -6694,6 +6869,27 @@ function trunc(x) {
 	return idiv(x, 1);
 }
 /**
+* @function
+* Round to the given number of decimals.
+*
+* Uses banker's rounding (round half to even) to match NumPy/JAX behavior.
+*/
+const round = jit$1(function round$1(a, decimals = 0) {
+	if (decimals === 0) return rint(a);
+	const factor = 10 ** decimals;
+	return rint(a.mul(factor)).mul(1 / factor);
+}, { staticArgnums: [1] });
+/**
+* @function
+* Round to the nearest integer, with ties going to the nearest even integer.
+*/
+const rint = jit$1(function rint$1(x) {
+	const rounded = floor(x.ref.add(.5));
+	const half = x.ref.sub(floor(x)).equal(.5);
+	const odd = remainder(rounded.ref, 2).notEqual(0);
+	return where(half.mul(odd), rounded.ref.sub(1), rounded);
+});
+/**
 * Compute `x1 * 2 ** x2` as a standard multiplication and exponentiation.
 *
 * This is the inverse of `frexp()`.
@@ -7021,6 +7217,7 @@ function triangularSolve(a, b, { leftSide = false, lower = false, transposeA = f
 //#region src/library/lax.ts
 var lax_exports = {};
 __export(lax_exports, {
+	bitcastConvertType: () => bitcastConvertType,
 	conv: () => conv,
 	convGeneralDilated: () => convGeneralDilated,
 	convTranspose: () => convTranspose,
@@ -7030,9 +7227,14 @@ __export(lax_exports, {
 	erfc: () => erfc,
 	linalg: () => lax_linalg_exports,
 	reduceWindow: () => reduceWindow,
-	stopGradient: () => stopGradient$1
+	stopGradient: () => stopGradient$1,
+	topK: () => topK
 });
 const JsArray = globalThis.Array;
+/** Elementwise bitcast an array into a new dtype. */
+function bitcastConvertType(x, newDtype) {
+	return fudgeArray(x).view(newDtype);
+}
 /**
 * General dot product/contraction operator.
 *
@@ -7254,6 +7456,39 @@ function erfc(x) {
 function stopGradient$1(x) {
 	return stopGradient(x);
 }
+/**
+* Returns top `k` values and their indices along the specified axis of operand.
+*
+* This is a _stable_ algorithm: If two elements are equal, the lower-index
+* element appears first.
+*
+* @returns A tuple of `(values, indices)`, where `values` and `indices` have
+* the same shape as `x`, except along `axis` where they have size `k`.
+*/
+function topK(x, k, axis = -1) {
+	x = fudgeArray(x);
+	axis = require_backend.checkAxis(axis, x.ndim);
+	const size$1 = x.shape[axis];
+	if (k < 0 || k > size$1) throw new Error(`topK: k must be in the range [0, ${size$1}], got ${k}`);
+	if (k === 0) {
+		const outShape = x.shape.slice();
+		outShape[axis] = 0;
+		const y$1 = zerosLike$1(x.ref, { shape: outShape });
+		const yi$1 = zerosLike$1(x, {
+			dtype: require_backend.DType.Int32,
+			shape: outShape
+		});
+		return [y$1, yi$1];
+	}
+	x = flip$1(x, [axis]);
+	x = moveaxis(x, axis, -1);
+	const [y, yi] = argsort$1(x);
+	const extract = (a) => {
+		a = a.slice(...require_backend.rep(a.ndim - 1, []), [-k]);
+		return flip$1(moveaxis(a, -1, axis), [axis]);
+	};
+	return [extract(y), extract(yi.neg().add(size$1 - 1))];
+}
 
 //#endregion
 //#region src/library/nn.ts
@@ -7445,7 +7680,7 @@ 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))));
+	} else return x.ref.mul(.5).mul(erfc$1(negative(x.mul(Math.SQRT1_2))));
 }, { staticArgnums: [1] });
 /**
 * Gated linear unit (GLU) activation function.
@@ -7703,6 +7938,7 @@ var random_exports = {};
 __export(random_exports, {
 	bernoulli: () => bernoulli,
 	bits: () => bits,
+	categorical: () => categorical,
 	cauchy: () => cauchy,
 	exponential: () => exponential,
 	gumbel: () => gumbel,
@@ -7730,7 +7966,9 @@ function getK01(key$1) {
 function key(seed) {
 	seed = array(seed, { dtype: require_backend.DType.Uint32 });
 	if (seed.ndim !== 0) throw new Error(`key: seed must be a scalar integer, but got shape ${seed.shape} - use jax.vmap for batching.`);
-	return stack([0, seed]);
+	const key$1 = stack([0, seed]);
+	if (key$1 instanceof Array$1) key$1._realizeSource();
+	return key$1;
 }
 /** Splits a PRNG key into `num` new keys by adding a leading axis. */
 function split(key$1, num = 2) {
@@ -7774,6 +8012,47 @@ function bernoulli(key$1, p = .5, shape$1 = []) {
 }
 /**
 * @function
+* Sample random values from categorical distributions.
+*
+* Uses the Gumbel max trick for sampling with replacement, or the Gumbel top-k
+* trick for sampling without replacement.
+*
+* Note: Sampling without replacement currently uses argsort and slices the last
+* k elements. This should be replaced with a more efficient topK implementation.
+*
+* - `key` - PRNG key
+* - `logits` - Unnormalized log probabilities of the categorical distribution(s).
+*   `softmax(logits, axis)` gives the corresponding probabilities.
+* - `axis` - Axis along which logits belong to the same categorical distribution.
+* - `shape` - Result batch shape. Must be broadcast-compatible with
+*   `logits.shape` with `axis` removed. Default is `logits.shape` with `axis` removed.
+* - `replace` - If true (default), sample with replacement. If false, sample
+*   without replacement (each category can only be selected once per batch).
+* @returns A random array with int dtype and shape given by `shape` if provided,
+*   otherwise `logits.shape` with `axis` removed.
+*/
+const categorical = jit$1(function categorical$1(key$1, logits, { axis = -1, shape: shape$1, replace = true } = {}) {
+	logits = fudgeArray(logits);
+	axis = require_backend.checkAxis(axis, logits.ndim);
+	const numCategories = logits.shape[axis];
+	const batchShape = logits.shape.toSpliced(axis, 1);
+	if (shape$1 === void 0) shape$1 = batchShape;
+	else if (!require_backend.deepEqual(require_backend.generalBroadcast(shape$1, batchShape), shape$1)) throw new Error(`Shape ${shape$1} is not broadcast-compatible with batch shape ${batchShape}.`);
+	const shapePrefix = shape$1.slice(0, shape$1.length - batchShape.length);
+	if (replace) {
+		const noise = gumbel(key$1, [...shapePrefix, ...logits.shape]);
+		return argmax(noise.add(logits), axis + shapePrefix.length);
+	} else {
+		const k = shapePrefix.reduce((a, b) => a * b, 1);
+		if (k > numCategories) throw new Error(`Number of samples without replacement (${k}) cannot exceed number of categories (${numCategories}).`);
+		const noise = gumbel(key$1, logits.shape);
+		const [values, indices] = topK(noise.add(logits), k, axis);
+		values.dispose();
+		return indices.reshape(shape$1);
+	}
+}, { staticArgnums: [2] });
+/**
+* @function
 * Sample from a Cauchy distribution with location 0 and scale 1.
 *
 * Uses inverse transform sampling: `x = tan(π * (u - 0.5))` where u ~ Uniform(0, 1).
@@ -7884,6 +8163,11 @@ Symbol.asyncDispose ??= Symbol.for("Symbol.asyncDispose");
 
 //#endregion
 //#region src/index.ts
+/** @namespace */
+const profiler = {
+	startTrace: require_backend.startTrace,
+	stopTrace: require_backend.stopTrace
+};
 /**
 * @function
 * Compute the forward-mode Jacobian-vector product for a function.
@@ -8080,6 +8364,7 @@ Object.defineProperty(exports, 'numpy', {
     return numpy_exports;
   }
 });
+exports.profiler = profiler;
 Object.defineProperty(exports, 'random', {
   enumerable: true,
   get: function () {