@jax-js/jax 0.1.9 → 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-DpI0riom.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));
@@ -1659,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)];
 	},
@@ -3081,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 {
@@ -4179,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 });
@@ -5254,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,
@@ -5284,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);
@@ -5299,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.
@@ -5356,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);
@@ -5539,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,
@@ -5556,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,
@@ -5596,7 +5623,6 @@ __export(numpy_exports, {
 	fullLike: () => fullLike$1,
 	greater: () => greater,
 	greaterEqual: () => greaterEqual,
-	hamming: () => hamming,
 	hann: () => hann,
 	heaviside: () => heaviside,
 	hstack: () => hstack,
@@ -5620,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,
@@ -5655,6 +5686,8 @@ __export(numpy_exports, {
 	remainder: () => remainder,
 	repeat: () => repeat,
 	reshape: () => reshape,
+	rint: () => rint,
+	round: () => round,
 	shape: () => shape,
 	sign: () => sign,
 	sin: () => sin,
@@ -5687,6 +5720,7 @@ __export(numpy_exports, {
 	var_: () => var_,
 	vdot: () => vdot,
 	vecdot: () => vecdot,
+	vecmat: () => vecmat,
 	vstack: () => vstack,
 	where: () => where,
 	zeros: () => zeros,
@@ -5750,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;
 /**
@@ -5857,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
@@ -6260,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");
@@ -6287,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);
@@ -6445,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));
@@ -6541,16 +6715,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.
 *
@@ -6696,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()`.
@@ -7023,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,
@@ -7036,6 +7231,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.
 *
@@ -7767,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) {
@@ -7962,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.
@@ -8158,6 +8364,7 @@ Object.defineProperty(exports, 'numpy', {
     return numpy_exports;
   }
 });
+exports.profiler = profiler;
 Object.defineProperty(exports, 'random', {
   enumerable: true,
   get: function () {

package/dist/index.d.cts

@@ -1004,6 +1004,8 @@ declare abstract class Tracer {
   reshape(shape: number | number[]): this;
   /** Copy the array and cast to a specified dtype. */
   astype(dtype: DType): this;
+  /** Return a bitwise cast of the array, viewed as a new dtype. */
+  view(dtype?: DType): this;
   /** Subtract an array from this one. */
   sub(other: this | TracerValue): this;
   /** Divide an array by this one. */
@@ -1427,8 +1429,10 @@ declare function triangularSolve(a: ArrayLike, b: ArrayLike, {
   unitDiagonal?: boolean;
 }): Array;
 declare namespace lax_d_exports {
-  export { DotDimensionNumbers, PaddingType, conv, convGeneralDilated, convTranspose, convWithGeneralPadding, dot$1 as dot, erf, erfc, lax_linalg_d_exports as linalg, reduceWindow, stopGradient, topK };
+  export { DotDimensionNumbers, PaddingType, bitcastConvertType, conv, convGeneralDilated, convTranspose, convWithGeneralPadding, dot$1 as dot, erf, erfc, lax_linalg_d_exports as linalg, reduceWindow, stopGradient, topK };
 }
+/** Elementwise bitcast an array into a new dtype. */
+declare function bitcastConvertType(x: ArrayLike, newDtype: DType): Array;
 /**
  * Dimension numbers for general `dot()` primitive.
  *
@@ -1567,7 +1571,7 @@ declare function fft(a: ComplexPair, axis?: number): ComplexPair;
  */
 declare function ifft(a: ComplexPair, axis?: number): ComplexPair;
 declare namespace numpy_linalg_d_exports {
-  export { cholesky, det, diagonal, inv, lstsq, matmul, matrixPower, matrixTranspose, outer, slogdet, solve, tensordot, trace, vecdot };
+  export { cholesky, cross$1 as cross, det, diagonal, inv, lstsq, matmul, matrixPower, matrixTranspose, outer, slogdet, solve, tensordot, trace, vecdot };
 }
 /**
  * Compute the Cholesky decomposition of a (batched) positive-definite matrix.
@@ -1582,6 +1586,13 @@ declare function cholesky(a: ArrayLike, {
   upper?: boolean;
   symmetrizeInput?: boolean;
 }): Array;
+/**
+ * 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.
+ */
+declare function cross$1(x1: ArrayLike, x2: ArrayLike, axis?: number): Array;
 /** Compute the determinant of a square matrix (batched). */
 declare function det(a: ArrayLike): Array;
 /** Compute the inverse of a square matrix (batched). */
@@ -1668,7 +1679,7 @@ type IInfo = Readonly<{
 /** Machine limits for integer types. */
 declare function iinfo(dtype: DType): IInfo;
 declare namespace numpy_d_exports {
-  export { Array, ArrayLike, DType, absolute as abs, absolute, acos, arccosh as acosh, add, all, allclose, any, arange, acos as arccos, arccosh, asin as arcsin, arcsinh, atan as arctan, atan2 as arctan2, arctanh, argmax, argmin, argsort, array, asin, arcsinh as asinh, astype, atan, atan2, arctanh as atanh, bool, broadcastArrays, broadcastShapes, broadcastTo, cbrt, ceil, clip, columnStack, concatenate, convolve, corrcoef, correlate, cos, cosh, cov, cumsum, cumsum as cumulativeSum, deg2rad, degrees, diag, diagonal, trueDivide as divide, divmod, dot, dstack, e, einsum, equal, eulerGamma, exp, exp2, expandDims, expm1, eye, numpy_fft_d_exports as fft, finfo, flip, fliplr, flipud, float16, float32, float64, floor, floorDivide, fmod, frexp, full, fullLike, greater, greaterEqual, hamming, hann, heaviside, hstack, hypot, identity$1 as identity, iinfo, inf, inner, int32, isfinite, isinf, isnan, isneginf, isposinf, ldexp, less, lessEqual, numpy_linalg_d_exports as linalg, linspace, log, log10, log1p, log2, logspace, matmul, matrixTranspose, max, maximum, mean, meshgrid, min, minimum, moveaxis, multiply, nan, nanToNum, ndim, negative, notEqual, ones, onesLike, outer, pad, transpose as permuteDims, pi, positive, power as pow, power, prod, promoteTypes, ptp, rad2deg, radians, ravel, reciprocal, remainder, repeat, reshape, shape$1 as shape, sign, sin, sinc, sinh, size, sort, split$1 as split, sqrt, square, squeeze, stack, std, subtract, sum, swapaxes, take, tan, tanh, tensordot, tile, trace, transpose, tri, tril, triu, trueDivide, trunc, uint32, var_, vdot, vecdot, vstack, where, zeros, zerosLike };
+  export { Array, ArrayLike, DType, absolute as abs, absolute, acos, arccosh as acosh, add, all, allclose, any, arange, acos as arccos, arccosh, asin as arcsin, arcsinh, atan as arctan, atan2 as arctan2, arctanh, argmax, argmin, argsort, round as around, array, arrayEqual, arrayEquiv, asin, arcsinh as asinh, astype, atan, atan2, arctanh as atanh, average, bool, broadcastArrays, broadcastShapes, broadcastTo, cbrt, ceil, clip, columnStack, concatenate, convolve, copysign, corrcoef, correlate, cos, cosh, cov, cross, cumsum, cumsum as cumulativeSum, deg2rad, degrees, diag, diagonal, trueDivide as divide, divmod, dot, dstack, e, einsum, equal, eulerGamma, exp, exp2, expandDims, expm1, eye, numpy_fft_d_exports as fft, finfo, flip, fliplr, flipud, float16, float32, float64, floor, floorDivide, fmod, frexp, full, fullLike, greater, greaterEqual, hann, heaviside, hstack, hypot, identity$1 as identity, iinfo, inf, inner, int32, isfinite, isinf, isnan, isneginf, isposinf, ldexp, less, lessEqual, numpy_linalg_d_exports as linalg, linspace, log, log10, log1p, log2, logicalAnd, logicalNot, logicalOr, logicalXor, logspace, matmul, matrixTranspose, matvec, max, maximum, mean, meshgrid, min, minimum, moveaxis, multiply, nan, nanToNum, ndim, negative, notEqual, ones, onesLike, outer, pad, transpose as permuteDims, pi, positive, power as pow, power, prod, promoteTypes, ptp, rad2deg, radians, ravel, reciprocal, remainder, repeat, reshape, rint, round, shape$1 as shape, sign, sin, sinc, sinh, size, sort, split$1 as split, sqrt, square, squeeze, stack, std, subtract, sum, swapaxes, take, tan, tanh, tensordot, tile, trace, transpose, tri, tril, triu, trueDivide, trunc, uint32, var_, vdot, vecdot, vecmat, vstack, where, zeros, zerosLike };
 }
 declare const float32 = DType.Float32;
 declare const int32 = DType.Int32;
@@ -1728,6 +1739,14 @@ declare const notEqual: (x: ArrayLike, y: ArrayLike) => Array;
 declare const greaterEqual: (x: ArrayLike, y: ArrayLike) => Array;
 /** @function Compare two arrays element-wise. */
 declare const lessEqual: (x: ArrayLike, y: ArrayLike) => Array;
+/** Compute element-wise logical AND. */
+declare function logicalAnd(x: ArrayLike, y: ArrayLike): Array;
+/** Compute element-wise logical OR. */
+declare function logicalOr(x: ArrayLike, y: ArrayLike): Array;
+/** Compute element-wise logical XOR. */
+declare function logicalXor(x: ArrayLike, y: ArrayLike): Array;
+/** Compute element-wise logical NOT. */
+declare function logicalNot(x: ArrayLike): Array;
 /** @function Element-wise ternary operator, evaluates to `x` if cond else `y`. */
 declare const where: (cond: ArrayLike, x: ArrayLike, y: ArrayLike) => Array;
 /**
@@ -1812,6 +1831,16 @@ declare function all(a: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
 declare function ptp(a: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
 /** Compute the average of the array elements along the specified axis. */
 declare function mean(a: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
+/**
+ * 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.
+ */
+declare function average(a: ArrayLike, axis?: Axis, opts?: {
+  weights?: ArrayLike;
+} & ReduceOpts): Array;
 /**
  * Returns the indices of the minimum values along an axis.
  *
@@ -1983,13 +2012,39 @@ declare function argsort(a: ArrayLike, axis?: number): Array;
  * numbered axis. By default, the flattened array is used.
  */
 declare function take(a: ArrayLike, indices: ArrayLike, axis?: number | null): Array;
-/** 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.
+ */
 declare function allclose(actual: Parameters<typeof array>[0], expected: Parameters<typeof array>[0], options?: {
   rtol?: number;
   atol?: number;
+  equalNaN?: boolean;
 }): boolean;
+/**
+ * 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.
+ */
+declare function arrayEqual(a1: ArrayLike, a2: ArrayLike, opts?: {
+  equalNaN?: boolean;
+}): Array;
+/**
+ * Check if two arrays are element-wise equal after broadcasting.
+ *
+ * Unlike `arrayEqual`, this allows inputs with different but
+ * broadcast-compatible shapes.
+ */
+declare function arrayEquiv(a1: ArrayLike, a2: ArrayLike): Array;
 /** Matrix product of two arrays. */
 declare function matmul(x: ArrayLike, y: ArrayLike): Array;
+/** Matrix-vector product. x1 is [..., M, N], x2 is [..., N] → [..., M]. */
+declare function matvec(x1: ArrayLike, x2: ArrayLike): Array;
+/** Vector-matrix product. x1 is [..., N], x2 is [..., N, M] → [..., M]. */
+declare function vecmat(x1: ArrayLike, x2: ArrayLike): Array;
 /** Dot product of two arrays. */
 declare function dot(x: ArrayLike, y: ArrayLike): Array;
 /**
@@ -2042,6 +2097,18 @@ declare function inner(x: ArrayLike, y: ArrayLike): Array;
  * be of shape `[x.size, y.size]`.
  */
 declare function outer(x: ArrayLike, y: ArrayLike): Array;
+/**
+ * @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`.
+ */
+declare const cross: OwnedFunction<(a: ArrayLike, b: ArrayLike, args_2?: {
+  axisa?: number | undefined;
+  axisb?: number | undefined;
+  axisc?: number | undefined;
+  axis?: number | undefined;
+} | undefined) => Array>;
 /** Vector dot product of two arrays along a given axis. */
 declare function vecdot(x: ArrayLike, y: ArrayLike, {
   axis
@@ -2087,14 +2154,13 @@ declare function clip(a: ArrayLike, min?: ArrayLike, max?: ArrayLike): Array;
 declare function absolute(x: ArrayLike): Array;
 /** Return an element-wise indication of sign of the input. */
 declare function sign(x: ArrayLike): Array;
-/** @function Return element-wise positive values of the input (no-op). */
-declare const positive: (x: ArrayLike) => Array;
 /**
- * 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.
  */
-declare function hamming(M: number): Array;
+declare const copysign: OwnedFunction<(x: ArrayLike, y: ArrayLike) => Array>;
+/** @function Return element-wise positive values of the input (no-op). */
+declare const positive: (x: ArrayLike) => Array;
 /**
  * Return the Hann window of size M, a taper with a weighted cosine bell.
  *
@@ -2189,6 +2255,18 @@ declare const remainder: OwnedFunction<(x: ArrayLike, y: ArrayLike) => Array>;
 declare function divmod(x: ArrayLike, y: ArrayLike): [Array, Array];
 /** Round input to the nearest integer towards zero. */
 declare function trunc(x: ArrayLike): Array;
+/**
+ * @function
+ * Round to the given number of decimals.
+ *
+ * Uses banker's rounding (round half to even) to match NumPy/JAX behavior.
+ */
+declare const round: OwnedFunction<(a: ArrayLike, decimals?: number | undefined) => Array>;
+/**
+ * @function
+ * Round to the nearest integer, with ties going to the nearest even integer.
+ */
+declare const rint: OwnedFunction<(x: ArrayLike) => Array>;
 /**
  * Compute `x1 * 2 ** x2` as a standard multiplication and exponentiation.
  *
@@ -2691,8 +2769,31 @@ declare namespace scipy_special_d_exports {
  * The logit function, `logit(p) = log(p / (1-p))`.
  */
 declare const logit: OwnedFunction<(x: ArrayLike) => Array>;
+//#endregion
+//#region src/tracing.d.ts
+/**
+ * Start collecting kernel traces.
+ *
+ * Traces appear in developer tools under the "Performance" tab, and they are
+ * useful for measuring fine-grained kernel execution time.
+ */
+declare function startTrace(): void;
+/**
+ * Stop collecting kernel traces.
+ *
+ * Traces appear in developer tools under the "Performance" tab, and they are
+ * useful for measuring fine-grained kernel execution time.
+ */
+declare function stopTrace(): void;
+/** Check if tracing is currently enabled. */
+
 //#endregion
 //#region src/index.d.ts
+/** @namespace */
+declare const profiler: {
+  startTrace: typeof startTrace;
+  stopTrace: typeof stopTrace;
+};
 /**
  * @function
  * Compute the forward-mode Jacobian-vector product for a function.
@@ -2857,4 +2958,4 @@ declare function blockUntilReady<T extends JsTree<any>>(x: T): Promise<T>;
  */
 declare function devicePut<T extends JsTree<any>>(x: T, device?: Device): Promise<MapJsTree<T, number | boolean, Array>>;
 //#endregion
-export { Array, ClosedJaxpr, DType, type Device, Jaxpr, type JsTree, type JsTreeDef, type OwnedFunction, blockUntilReady, defaultDevice, devicePut, devices, grad, hessian, init, jacfwd, jacrev as jacobian, jacrev, jit, jvp, lax_d_exports as lax, linearize, makeJaxpr, nn_d_exports as nn, numpy_d_exports as numpy, random_d_exports as random, scipy_special_d_exports as scipySpecial, setDebug, tree_d_exports as tree, valueAndGrad, vjp, vmap };
+export { Array, ClosedJaxpr, DType, type Device, Jaxpr, type JsTree, type JsTreeDef, type OwnedFunction, blockUntilReady, defaultDevice, devicePut, devices, grad, hessian, init, jacfwd, jacrev as jacobian, jacrev, jit, jvp, lax_d_exports as lax, linearize, makeJaxpr, nn_d_exports as nn, numpy_d_exports as numpy, profiler, random_d_exports as random, scipy_special_d_exports as scipySpecial, setDebug, tree_d_exports as tree, valueAndGrad, vjp, vmap };