@jax-js/jax 0.1.9 → 0.1.11

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, Routine, Routines, ShapeTracker, accessorAluExp, accessorGlobal, assertNonNull, byteWidth, checkAxis, checkInts, 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-BId79r5b.js";
+import { AluExp, AluGroup, AluOp, AluVar, DEBUG, DType, FpHash, Kernel, PPrint, Reduction, Routine, Routines, ShapeTracker, accessorAluExp, accessorGlobal, assertNonNull, byteWidth, checkAxis, checkInts, deepEqual, defaultDevice, devices, dtypedArray, dtypedJsArray, generalBroadcast, getBackend, getWebGPUDevice, init, invertPermutation, isFloatDtype, isNumberPair, isPermutation, normalizeAxis, partitionList, prod, promoteTypes, range, recursiveFlatten, rep, runWithCache, setDebug, startTrace, stopTrace, toposort, unravelAlu, unzip2, zip, zipn } from "./backend-DZvR7mZV.js";
 
 //#region src/frontend/convolution.ts
 /**
@@ -333,6 +333,8 @@ let Primitive = /* @__PURE__ */ function(Primitive$1) {
 	Primitive$1["Mod"] = "mod";
 	Primitive$1["Min"] = "min";
 	Primitive$1["Max"] = "max";
+	Primitive$1["BitCombine"] = "bit_combine";
+	Primitive$1["BitShift"] = "bit_shift";
 	Primitive$1["Neg"] = "neg";
 	Primitive$1["Reciprocal"] = "reciprocal";
 	Primitive$1["Floor"] = "floor";
@@ -406,6 +408,12 @@ function min$1(x, y) {
 function max$1(x, y) {
 	return bind1(Primitive.Max, [x, y]);
 }
+function bitCombine(x, y, op) {
+	return bind1(Primitive.BitCombine, [x, y], { op });
+}
+function bitShift(x, y, op) {
+	return bind1(Primitive.BitShift, [x, y], { op });
+}
 function neg(x) {
 	return bind1(Primitive.Neg, [x]);
 }
@@ -807,6 +815,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));
@@ -1615,6 +1628,16 @@ const abstractEvalRules = {
 	[Primitive.Mod]: binopAbstractEval,
 	[Primitive.Min]: binopAbstractEval,
 	[Primitive.Max]: binopAbstractEval,
+	[Primitive.BitCombine]([x, y]) {
+		const aval = promoteAvals(x, y);
+		if (isFloatDtype(aval.dtype)) throw new TypeError(`bitwise operations require integer or boolean inputs, got ${aval.dtype}`);
+		return [aval];
+	},
+	[Primitive.BitShift]([x, y]) {
+		const shape$1 = generalBroadcast(x.shape, y.shape);
+		if (isFloatDtype(x.dtype) || isFloatDtype(y.dtype) || x.dtype === DType.Bool || y.dtype === DType.Bool) throw new TypeError(`bit shift operations require integer inputs, got ${x} and ${y}`);
+		return [new ShapedArray(shape$1, x.dtype, x.weakType)];
+	},
 	[Primitive.Neg]: vectorizedUnopAbstractEval,
 	[Primitive.Reciprocal]: vectorizedUnopAbstractEval,
 	[Primitive.Floor]: vectorizedUnopAbstractEval,
@@ -1624,7 +1647,7 @@ const abstractEvalRules = {
 		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 (x.dtype !== dtype && (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, false)];
 	},
@@ -2150,6 +2173,8 @@ const jitRules = {
 	[Primitive.Mod]: broadcastedJit(([a, b]) => AluExp.mod(a, b)),
 	[Primitive.Min]: broadcastedJit(([a, b]) => AluExp.min(a, b)),
 	[Primitive.Max]: broadcastedJit(([a, b]) => AluExp.max(a, b)),
+	[Primitive.BitCombine]: broadcastedJit(([a, b], { op }) => AluExp.bitCombine(a, b, op)),
+	[Primitive.BitShift]: broadcastedJit(([a, b], { op }) => AluExp.bitShift(a, b, op)),
 	[Primitive.Neg]: unopJit((a) => AluExp.sub(AluExp.const(a.dtype, 0), a)),
 	[Primitive.Reciprocal]: unopJit(AluExp.reciprocal),
 	[Primitive.Floor]: unopJit(AluExp.floor),
@@ -2342,7 +2367,9 @@ function splitGraphDataflow(backend, jaxpr) {
 					case Primitive.Idiv:
 					case Primitive.Mod:
 					case Primitive.Min:
-					case Primitive.Max: {
+					case Primitive.Max:
+					case Primitive.BitCombine:
+					case Primitive.BitShift: {
 						const otherInput = nextEqn.inputs.find((v) => v !== outVar);
 						if (otherInput instanceof Lit || deepEqual(generalBroadcast(otherInput.aval.shape, outVar.aval.shape), outVar.aval.shape)) {
 							head = usages[0];
@@ -2981,6 +3008,42 @@ var Array$1 = class Array$1 extends Tracer {
 		return dtypedArray(this.dtype, buf);
 	}
 	/**
+	* Return this array as a WebGPU buffer (with `STORAGE | COPY_SRC`).
+	*
+	* Only available on the WebGPU backend. The array's memory is still managed
+	* by jax-js, and it will be freed when the buffer is no longer in use. You
+	* _should not_ mutate the buffer's contents.
+	*
+	* Note that the GPU buffer may be slightly larger than the array's size; it
+	* will always be aligned to 4 bytes.
+	*/
+	async gpuBuffer() {
+		if (this.device !== "webgpu") throw new Error(`gpuBuffer() is only available on WebGPU backend`);
+		this.#realize();
+		const pending = this.#pending;
+		if (pending) {
+			await Promise.all(pending.map((p) => p.prepare()));
+			for (const p of pending) p.submit();
+		}
+		const backend = this.#backend;
+		const { buffer } = backend.buffers.get(this.#source);
+		this.dispose();
+		return buffer;
+	}
+	/** Synchronous version of `Array.gpuBuffer()`. */
+	gpuBufferSync() {
+		if (this.device !== "webgpu") throw new Error(`gpuBufferSync() is only available on WebGPU backend`);
+		this.#realize();
+		for (const p of this.#pending) {
+			p.prepareSync();
+			p.submit();
+		}
+		const backend = this.#backend;
+		const { buffer } = backend.buffers.get(this.#source);
+		this.dispose();
+		return buffer;
+	}
+	/**
 	* Convert this array into a JavaScript object.
 	*
 	* This is a blocking operation that will compile all of the shaders and wait
@@ -3027,6 +3090,14 @@ var Array$1 = class Array$1 extends Tracer {
 			[Primitive.Max]([x, y]) {
 				return [x.#binary(AluOp.Max, y)];
 			},
+			[Primitive.BitCombine]([x, y], { op }) {
+				const custom = (src) => AluExp.bitCombine(src[0], src[1], op);
+				return [Array$1.#naryCustom("bit_combine", custom, [x, y])];
+			},
+			[Primitive.BitShift]([x, y], { op }) {
+				const custom = (src) => AluExp.bitShift(src[0], src[1], op);
+				return [Array$1.#naryCustom("bit_shift", custom, [x, y], { dtypeOverride: [void 0, y.dtype] })];
+			},
 			[Primitive.Neg]([x]) {
 				return [zerosLike$1(x.ref).#binary(AluOp.Sub, x)];
 			},
@@ -3046,8 +3117,8 @@ var Array$1 = class Array$1 extends Tracer {
 				return [x.#unary(AluOp.Cast, dtype)];
 			},
 			[Primitive.Bitcast]([x], { dtype }) {
-				if (x.dtype === DType.Bool || dtype === DType.Bool) throw new TypeError("Bitcast to/from bool is not allowed");
 				if (x.dtype === dtype) return [x];
+				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`);
 				if (x.#source instanceof AluExp) return [x.#unary(AluOp.Bitcast, dtype)];
 				else {
@@ -3718,6 +3789,8 @@ const vmapRules = {
 	[Primitive.Mod]: broadcastBatcher(Primitive.Mod),
 	[Primitive.Min]: broadcastBatcher(Primitive.Min),
 	[Primitive.Max]: broadcastBatcher(Primitive.Max),
+	[Primitive.BitCombine]: broadcastBatcher(Primitive.BitCombine),
+	[Primitive.BitShift]: broadcastBatcher(Primitive.BitShift),
 	[Primitive.Neg]: unopBatcher(Primitive.Neg),
 	[Primitive.Reciprocal]: unopBatcher(Primitive.Reciprocal),
 	[Primitive.Floor]: unopBatcher(Primitive.Floor),
@@ -4040,6 +4113,8 @@ const jvpRules = {
 	[Primitive.Max]([x, y], [dx, dy]) {
 		return [[max$1(x.ref, y.ref)], [where$1(less$1(x, y), dy, dx)]];
 	},
+	[Primitive.BitCombine]: zeroTangentsJvp(Primitive.BitCombine),
+	[Primitive.BitShift]: zeroTangentsJvp(Primitive.BitShift),
 	[Primitive.Neg]: linearTangentsJvp(Primitive.Neg),
 	[Primitive.Reciprocal]([x], [dx]) {
 		const xRecip = reciprocal$1(x.ref);
@@ -4142,6 +4217,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 });
@@ -5217,6 +5293,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,
@@ -5229,7 +5306,8 @@ __export(numpy_linalg_exports, {
 	solve: () => solve,
 	tensordot: () => tensordot,
 	trace: () => trace,
-	vecdot: () => vecdot
+	vecdot: () => vecdot,
+	vectorNorm: () => vectorNorm
 });
 function checkSquare(name, a) {
 	if (a.ndim < 2 || a.shape[a.ndim - 1] !== a.shape[a.ndim - 2]) throw new Error(`${name}: input must be at least 2D square matrix, got ${a.aval}`);
@@ -5247,6 +5325,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 = checkAxis(axis, ndim(x1));
+	const a2 = 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);
@@ -5262,7 +5353,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.
@@ -5319,8 +5410,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);
@@ -5394,6 +5486,23 @@ function solve(a, b) {
 	if (bIs1d) x = squeeze(x, -1);
 	return x;
 }
+/**
+* Compute the vector norm of an array.
+*
+* @param x - Input array.
+* @param ord - Order of the norm (default 2). Supports `Infinity`, `-Infinity`, `0`, or any real number.
+* @param axis - Axis/axes to reduce over (default: all axes).
+* @param keepdims - Whether to keep reduced dimensions as size 1.
+* @returns The norm of `x`, reduced over the given axes.
+*/
+function vectorNorm(x, { ord = 2, axis = null, keepdims = false } = {}) {
+	x = fudgeArray(x);
+	const ax = axis ?? null;
+	if (ord === Infinity) return max(absolute(x), ax, { keepdims });
+	else if (ord === -Infinity) return min(absolute(x), ax, { keepdims });
+	else if (ord === 0) return x.notEqual(0).astype(x.dtype).sum(ax, { keepdims });
+	else return power(power(absolute(x), ord).sum(ax, { keepdims }), 1 / ord);
+}
 
 //#endregion
 //#region src/library/numpy/dtype-info.ts
@@ -5502,13 +5611,24 @@ __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,
+	bitwiseAnd: () => bitwiseAnd,
+	bitwiseInvert: () => invert,
+	bitwiseLeftShift: () => leftShift,
+	bitwiseNot: () => invert,
+	bitwiseOr: () => bitwiseOr,
+	bitwiseRightShift: () => rightShift,
+	bitwiseXor: () => bitwiseXor,
 	bool: () => bool,
 	broadcastArrays: () => broadcastArrays,
 	broadcastShapes: () => broadcastShapes,
@@ -5519,11 +5639,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,
@@ -5559,7 +5681,6 @@ __export(numpy_exports, {
 	fullLike: () => fullLike$1,
 	greater: () => greater,
 	greaterEqual: () => greaterEqual,
-	hamming: () => hamming,
 	hann: () => hann,
 	heaviside: () => heaviside,
 	hstack: () => hstack,
@@ -5569,12 +5690,14 @@ __export(numpy_exports, {
 	inf: () => inf,
 	inner: () => inner,
 	int32: () => int32,
+	invert: () => invert,
 	isfinite: () => isfinite,
 	isinf: () => isinf,
 	isnan: () => isnan,
 	isneginf: () => isneginf,
 	isposinf: () => isposinf,
 	ldexp: () => ldexp,
+	leftShift: () => leftShift,
 	less: () => less,
 	lessEqual: () => lessEqual,
 	linalg: () => numpy_linalg_exports,
@@ -5583,9 +5706,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,
@@ -5618,6 +5746,9 @@ __export(numpy_exports, {
 	remainder: () => remainder,
 	repeat: () => repeat,
 	reshape: () => reshape,
+	rightShift: () => rightShift,
+	rint: () => rint,
+	round: () => round,
 	shape: () => shape,
 	sign: () => sign,
 	sin: () => sin,
@@ -5650,6 +5781,7 @@ __export(numpy_exports, {
 	var_: () => var_,
 	vdot: () => vdot,
 	vecdot: () => vecdot,
+	vecmat: () => vecmat,
 	vstack: () => vstack,
 	where: () => where,
 	zeros: () => zeros,
@@ -5713,6 +5845,60 @@ 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, DType.Bool).mul(astype(y, DType.Bool));
+}
+/** Compute element-wise logical OR. */
+function logicalOr(x, y) {
+	return astype(x, DType.Bool).add(astype(y, DType.Bool));
+}
+/** Compute element-wise logical XOR. */
+function logicalXor(x, y) {
+	return notEqual(astype(x, DType.Bool), astype(y, DType.Bool));
+}
+/** Compute element-wise logical NOT. */
+function logicalNot(x) {
+	return notEqual(astype(x, DType.Bool), true);
+}
+/** Compute element-wise bitwise AND. */
+function bitwiseAnd(x, y) {
+	return bitCombine(x, y, "and");
+}
+/** Compute element-wise bitwise OR. */
+function bitwiseOr(x, y) {
+	return bitCombine(x, y, "or");
+}
+/** Compute element-wise bitwise XOR. */
+function bitwiseXor(x, y) {
+	return bitCombine(x, y, "xor");
+}
+/** Compute element-wise bitwise NOT (inversion). */
+function invert(x) {
+	const arr = fudgeArray(x);
+	let allOnes;
+	switch (arr.dtype) {
+		case DType.Bool:
+			allOnes = true;
+			break;
+		case DType.Uint32:
+			allOnes = 4294967295;
+			break;
+		case DType.Int32:
+			allOnes = -1;
+			break;
+		default: throw new TypeError(`invert: unsupported dtype ${arr.dtype}`);
+	}
+	return bitCombine(arr, allOnes, "xor");
+}
+/** Compute element-wise left bit shift. */
+function leftShift(x, y) {
+	return bitShift(x, y, "shl");
+}
+/** Compute element-wise right bit shift. */
+function rightShift(x, y) {
+	return bitShift(x, y, "shr");
+}
 /** @function Element-wise ternary operator, evaluates to `x` if cond else `y`. */
 const where = where$1;
 /**
@@ -5820,6 +6006,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 = normalizeAxis(axis, ndim(a));
+	const wShape = weights.shape;
+	const aShape = a.shape;
+	if (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
@@ -6223,20 +6437,63 @@ function take(a, indices, axis = null) {
 	axis = 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 (!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 (!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");
@@ -6250,6 +6507,16 @@ function matmul(x, y) {
 		rhsBatchDims: 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);
@@ -6408,6 +6675,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 = checkAxis(axisa, ndim(a));
+	axisb = 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 = checkAxis(axis ?? -1, ndim(x));
@@ -6504,16 +6814,15 @@ function sign(x) {
 	x = fudgeArray(x);
 	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.
 *
@@ -6659,6 +6968,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()`.
@@ -6986,6 +7316,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,
@@ -6999,6 +7330,10 @@ __export(lax_exports, {
 	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.
 *
@@ -7730,7 +8065,9 @@ function getK01(key$1) {
 function key(seed) {
 	seed = array(seed, { dtype: 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) {
@@ -7925,6 +8262,11 @@ Symbol.asyncDispose ??= Symbol.for("Symbol.asyncDispose");
 
 //#endregion
 //#region src/index.ts
+/** @namespace */
+const profiler = {
+	startTrace,
+	stopTrace
+};
 /**
 * @function
 * Compute the forward-mode Jacobian-vector product for a function.
@@ -8085,4 +8427,4 @@ async function devicePut(x, device) {
 }
 
 //#endregion
-export { Array$1 as Array, ClosedJaxpr, DType, Jaxpr, blockUntilReady, defaultDevice, devicePut, devices, grad, hessian, init, jacfwd, jacrev as 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 };
+export { Array$1 as Array, ClosedJaxpr, DType, Jaxpr, blockUntilReady, defaultDevice, devicePut, devices, getWebGPUDevice, grad, hessian, init, jacfwd, jacrev as jacobian, jacrev, jit, jvp, lax_exports as lax, linearize, makeJaxpr, nn_exports as nn, numpy_exports as numpy, profiler, random_exports as random, scipy_special_exports as scipySpecial, setDebug, tree_exports as tree, valueAndGrad, vjp, vmap };

package/dist/{webgl-DnGrclTz.js → webgl-D8-14NzA.js}

@@ -1,4 +1,4 @@
-import { AluGroup, AluOp, DEBUG, DType, Executable, SlotError, UnsupportedOpError, UnsupportedRoutineError, isFloatDtype, range, strip1, tuneNullopt } from "./backend-BId79r5b.js";
+import { AluGroup, AluOp, DEBUG, DType, Executable, SlotError, UnsupportedOpError, UnsupportedRoutineError, isFloatDtype, range, strip1, tuneNullopt } from "./backend-DZvR7mZV.js";
 
 //#region src/backend/webgl/builtins.ts
 const threefrySrc = `
@@ -458,6 +458,12 @@ function generateExpression(exp, args, inputDtypes) {
 			else source = `min(${a}, ${b})`;
 			else if (op === AluOp.Max) if (dtype === DType.Bool) source = `(${a} || ${b})`;
 			else source = `max(${a}, ${b})`;
+			else if (op === AluOp.BitCombine) {
+				let infix = arg === "and" ? "&" : arg === "or" ? "|" : "^";
+				if (dtype === DType.Bool) infix = infix + infix;
+				source = `(${a} ${infix} ${b})`;
+			} else if (op === AluOp.BitShift) if (arg === "shl") source = `(${a} << ${b})`;
+			else source = `(${a} >> ${b})`;
 		} else if (AluGroup.Compare.has(op)) {
 			const a = gen(src[0]);
 			const b = gen(src[1]);