@jax-js/jax 0.1.9 → 0.1.11

package/dist/index.d.cts

@@ -232,6 +232,8 @@ declare class AluExp implements FpHashable {
   static cast(dtype: DType, a: AluExp): AluExp;
   static bitcast(dtype: DType, a: AluExp): AluExp;
   static threefry2x32(k0: AluExp, k1: AluExp, c0: AluExp, c1: AluExp, mode?: "xor" | 0 | 1): AluExp;
+  static bitCombine(a: AluExp, b: AluExp, mode: "and" | "or" | "xor"): AluExp;
+  static bitShift(a: AluExp, b: AluExp, mode: "shl" | "shr"): AluExp;
   static cmplt(a: AluExp, b: AluExp): AluExp;
   static cmpne(a: AluExp, b: AluExp): AluExp;
   static where(cond: AluExp, a: AluExp, b: AluExp): AluExp;
@@ -323,6 +325,11 @@ declare enum AluOp {
   Reciprocal = "Reciprocal",
   Cast = "Cast",
   Bitcast = "Bitcast",
+  BitCombine = "BitCombine",
+  // arg = 'or' | 'and' | 'xor'
+  BitInvert = "BitInvert",
+  BitShift = "BitShift",
+  // arg = 'shl' | 'shr'
   Cmplt = "Cmplt",
   Cmpne = "Cmpne",
   Where = "Where",
@@ -546,6 +553,11 @@ declare class Executable<T = any> {
   source: Kernel | Routine, /** Extra data specific to the backend running this executable. */
   data: T);
 }
+/**
+ * If the WebGPU backend has been initialized, return the `GPUDevice` that this
+ * backend runs on. This is useful for sharing buffers.
+ */
+declare function getWebGPUDevice(): GPUDevice;
 declare namespace tree_d_exports {
   export { JsTree, JsTreeDef, MapJsTree, NodeType, dispose, flatten, leaves, map, ref, structure, unflatten };
 }
@@ -719,6 +731,8 @@ declare enum Primitive {
   // uses sign of numerator, C-style, matches JS but not Python
   Min = "min",
   Max = "max",
+  BitCombine = "bit_combine",
+  BitShift = "bit_shift",
   Neg = "neg",
   Reciprocal = "reciprocal",
   Floor = "floor",
@@ -767,6 +781,12 @@ declare enum Primitive {
   Jit = "jit",
 }
 interface PrimitiveParamsImpl extends Record<Primitive, Record<string, any>> {
+  [Primitive.BitCombine]: {
+    op: "and" | "or" | "xor";
+  };
+  [Primitive.BitShift]: {
+    op: "shl" | "shr";
+  };
   [Primitive.Cast]: {
     dtype: DType;
   };
@@ -1004,6 +1024,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. */
@@ -1192,6 +1214,19 @@ declare class Array extends Tracer {
    * recommended for performance reasons, as it will block rendering.
    */
   dataSync(): DataArray;
+  /**
+   * 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.
+   */
+  gpuBuffer(): Promise<GPUBuffer>;
+  /** Synchronous version of `Array.gpuBuffer()`. */
+  gpuBufferSync(): GPUBuffer;
   /**
    * Convert this array into a JavaScript object.
    *
@@ -1427,8 +1462,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 +1604,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, vectorNorm };
 }
 /**
  * Compute the Cholesky decomposition of a (batched) positive-definite matrix.
@@ -1582,6 +1619,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). */
@@ -1615,6 +1659,24 @@ declare function slogdet(a: ArrayLike): [Array, Array];
  * @returns Solution `x` of shape `(..., N)` or `(..., N, M)`.
  */
 declare function solve(a: ArrayLike, b: ArrayLike): Array;
+/**
+ * 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.
+ */
+declare function vectorNorm(x: ArrayLike, {
+  ord,
+  axis,
+  keepdims
+}?: {
+  ord?: number;
+  axis?: number | number[] | null;
+  keepdims?: boolean;
+}): Array;
 //#endregion
 //#region src/library/numpy/dtype-info.d.ts
 /** @inline */
@@ -1668,7 +1730,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, bitwiseAnd, invert as bitwiseInvert, leftShift as bitwiseLeftShift, invert as bitwiseNot, bitwiseOr, rightShift as bitwiseRightShift, bitwiseXor, 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, invert, isfinite, isinf, isnan, isneginf, isposinf, ldexp, leftShift, 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, rightShift, 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 +1790,26 @@ 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;
+/** Compute element-wise bitwise AND. */
+declare function bitwiseAnd(x: ArrayLike, y: ArrayLike): Array;
+/** Compute element-wise bitwise OR. */
+declare function bitwiseOr(x: ArrayLike, y: ArrayLike): Array;
+/** Compute element-wise bitwise XOR. */
+declare function bitwiseXor(x: ArrayLike, y: ArrayLike): Array;
+/** Compute element-wise bitwise NOT (inversion). */
+declare function invert(x: ArrayLike): Array;
+/** Compute element-wise left bit shift. */
+declare function leftShift(x: ArrayLike, y: ArrayLike): Array;
+/** Compute element-wise right bit shift. */
+declare function rightShift(x: ArrayLike, y: 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 +1894,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 +2075,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 +2160,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 +2217,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 +2318,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 +2832,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 +3021,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, getWebGPUDevice, 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 };

package/dist/index.d.ts

@@ -229,6 +229,8 @@ declare class AluExp implements FpHashable {
   static cast(dtype: DType, a: AluExp): AluExp;
   static bitcast(dtype: DType, a: AluExp): AluExp;
   static threefry2x32(k0: AluExp, k1: AluExp, c0: AluExp, c1: AluExp, mode?: "xor" | 0 | 1): AluExp;
+  static bitCombine(a: AluExp, b: AluExp, mode: "and" | "or" | "xor"): AluExp;
+  static bitShift(a: AluExp, b: AluExp, mode: "shl" | "shr"): AluExp;
   static cmplt(a: AluExp, b: AluExp): AluExp;
   static cmpne(a: AluExp, b: AluExp): AluExp;
   static where(cond: AluExp, a: AluExp, b: AluExp): AluExp;
@@ -320,6 +322,11 @@ declare enum AluOp {
   Reciprocal = "Reciprocal",
   Cast = "Cast",
   Bitcast = "Bitcast",
+  BitCombine = "BitCombine",
+  // arg = 'or' | 'and' | 'xor'
+  BitInvert = "BitInvert",
+  BitShift = "BitShift",
+  // arg = 'shl' | 'shr'
   Cmplt = "Cmplt",
   Cmpne = "Cmpne",
   Where = "Where",
@@ -543,6 +550,11 @@ declare class Executable<T = any> {
   source: Kernel | Routine, /** Extra data specific to the backend running this executable. */
   data: T);
 }
+/**
+ * If the WebGPU backend has been initialized, return the `GPUDevice` that this
+ * backend runs on. This is useful for sharing buffers.
+ */
+declare function getWebGPUDevice(): GPUDevice;
 declare namespace tree_d_exports {
   export { JsTree, JsTreeDef, MapJsTree, NodeType, dispose, flatten, leaves, map, ref, structure, unflatten };
 }
@@ -716,6 +728,8 @@ declare enum Primitive {
   // uses sign of numerator, C-style, matches JS but not Python
   Min = "min",
   Max = "max",
+  BitCombine = "bit_combine",
+  BitShift = "bit_shift",
   Neg = "neg",
   Reciprocal = "reciprocal",
   Floor = "floor",
@@ -764,6 +778,12 @@ declare enum Primitive {
   Jit = "jit",
 }
 interface PrimitiveParamsImpl extends Record<Primitive, Record<string, any>> {
+  [Primitive.BitCombine]: {
+    op: "and" | "or" | "xor";
+  };
+  [Primitive.BitShift]: {
+    op: "shl" | "shr";
+  };
   [Primitive.Cast]: {
     dtype: DType;
   };
@@ -1001,6 +1021,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. */
@@ -1189,6 +1211,19 @@ declare class Array extends Tracer {
    * recommended for performance reasons, as it will block rendering.
    */
   dataSync(): DataArray;
+  /**
+   * 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.
+   */
+  gpuBuffer(): Promise<GPUBuffer>;
+  /** Synchronous version of `Array.gpuBuffer()`. */
+  gpuBufferSync(): GPUBuffer;
   /**
    * Convert this array into a JavaScript object.
    *
@@ -1424,8 +1459,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.
  *
@@ -1564,7 +1601,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, vectorNorm };
 }
 /**
  * Compute the Cholesky decomposition of a (batched) positive-definite matrix.
@@ -1579,6 +1616,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). */
@@ -1612,6 +1656,24 @@ declare function slogdet(a: ArrayLike): [Array, Array];
  * @returns Solution `x` of shape `(..., N)` or `(..., N, M)`.
  */
 declare function solve(a: ArrayLike, b: ArrayLike): Array;
+/**
+ * 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.
+ */
+declare function vectorNorm(x: ArrayLike, {
+  ord,
+  axis,
+  keepdims
+}?: {
+  ord?: number;
+  axis?: number | number[] | null;
+  keepdims?: boolean;
+}): Array;
 //#endregion
 //#region src/library/numpy/dtype-info.d.ts
 /** @inline */
@@ -1665,7 +1727,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, bitwiseAnd, invert as bitwiseInvert, leftShift as bitwiseLeftShift, invert as bitwiseNot, bitwiseOr, rightShift as bitwiseRightShift, bitwiseXor, 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, invert, isfinite, isinf, isnan, isneginf, isposinf, ldexp, leftShift, 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, rightShift, 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;
@@ -1725,6 +1787,26 @@ 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;
+/** Compute element-wise bitwise AND. */
+declare function bitwiseAnd(x: ArrayLike, y: ArrayLike): Array;
+/** Compute element-wise bitwise OR. */
+declare function bitwiseOr(x: ArrayLike, y: ArrayLike): Array;
+/** Compute element-wise bitwise XOR. */
+declare function bitwiseXor(x: ArrayLike, y: ArrayLike): Array;
+/** Compute element-wise bitwise NOT (inversion). */
+declare function invert(x: ArrayLike): Array;
+/** Compute element-wise left bit shift. */
+declare function leftShift(x: ArrayLike, y: ArrayLike): Array;
+/** Compute element-wise right bit shift. */
+declare function rightShift(x: ArrayLike, y: ArrayLike): Array;
 /** @function Element-wise ternary operator, evaluates to `x` if cond else `y`. */
 declare const where: (cond: ArrayLike, x: ArrayLike, y: ArrayLike) => Array;
 /**
@@ -1809,6 +1891,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.
  *
@@ -1980,13 +2072,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;
 /**
@@ -2039,6 +2157,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
@@ -2084,14 +2214,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.
  *
@@ -2186,6 +2315,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.
  *
@@ -2688,8 +2829,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.
@@ -2854,4 +3018,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, getWebGPUDevice, 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 };