@jax-js/jax 0.1.4 → 0.1.6

package/dist/index.d.ts

@@ -437,10 +437,37 @@ declare enum Routines {
   Sort = "Sort",
   /** Returns `int32` indices of the stably sorted array. */
   Argsort = "Argsort",
-  /** Solve a triangular system of questions. */
+  /**
+   * Solve a triangular system of equations.
+   *
+   * The first batch of inputs `A` should be of shape `[..., N, N]` and upper
+   * triangular, while the second batch `B` should be of shape `[..., M, N]`.
+   *
+   * Solves for `X` in the equation `A @ X.T = B.T`, where `A` is the
+   * triangular matrix. This is equivalent to `X = B @ A^-T`.
+   */
   TriangularSolve = "TriangularSolve",
-  /** Cholesky decomposition of 2D positive semi-definite matrices. */
+  /**
+   * Cholesky decomposition of 2D positive semi-definite matrices.
+   *
+   * The input batch should be of shape `[..., N, N]`, and the output batch is
+   * of the same shape, containing the lower-triangular matrix `L` such that
+   * `A = L @ L.T`. Behavior is unspecified if A is not positive semi-definite.
+   */
   Cholesky = "Cholesky",
+  /**
+   * LU decomposition of 2D rectangular matrices.
+   *
+   * The input is a batch of shape `[..., M, N]`, and the output is a tuple of
+   * three arrays: `LU, Pivots, Permutation`.
+   *
+   * - `LU` is of shape `[..., M, N]`, containing the combined lower and upper
+   *   triangular matrices. (lower triangular = implicit unit diagonal)
+   * - `Pivots` is of shape `[..., min(M, N)]`, containing the row swaps.
+   * - `Permutation` is of shape `[..., M]`, containing the permutation vector
+   *   such that `P = eye(M).slice(Permutation)` -> `P @ A = L @ U`.
+   */
+  LU = "LU",
 }
 interface RoutineType {
   inputShapes: number[][];
@@ -450,7 +477,7 @@ interface RoutineType {
 }
 //#endregion
 //#region src/backend.d.ts
-type Device = "cpu" | "wasm" | "webgpu";
+type Device = "cpu" | "wasm" | "webgpu" | "webgl";
 declare const devices: Device[];
 /** Configure the default device for arrays. */
 declare function defaultDevice(device?: Device): Device;
@@ -535,7 +562,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$1 as cholesky, diagonal, lstsq, matmul, matrixTranspose, outer, tensordot, trace, vecdot };
+  export { cholesky$1 as cholesky, det, diagonal, inv, lstsq, matmul, matrixPower, matrixTranspose, outer, slogdet, solve, tensordot, trace, vecdot };
 }
 /**
  * Compute the Cholesky decomposition of a (batched) positive-definite matrix.
@@ -550,6 +577,10 @@ declare function cholesky$1(a: ArrayLike, {
   upper?: boolean;
   symmetrizeInput?: boolean;
 }): Array;
+/** Compute the determinant of a square matrix (batched). */
+declare function det(a: ArrayLike): Array;
+/** Compute the inverse of a square matrix (batched). */
+declare function inv(a: ArrayLike): Array;
 /**
  * Return the least-squares solution to a linear equation.
  *
@@ -564,8 +595,75 @@ declare function cholesky$1(a: ArrayLike, {
  * @return least-squares solution of shape `(N,)` or `(N, K)`
  */
 declare function lstsq(a: ArrayLike, b: ArrayLike): Array;
+/** Raise a square matrix to an integer power, via repeated squarings. */
+declare function matrixPower(a: ArrayLike, n: number): Array;
+/** Return sign and natural logarithm of the determinant of `a`. */
+declare function slogdet(a: ArrayLike): [Array, Array];
+/**
+ * Solve a linear system of equations.
+ *
+ * This solves a (batched) linear system of equations `a @ x = b` for `x` given
+ * `a` and `b`. If `a` is singular, this will return `nan` or `inf` values.
+ *
+ * @param a - Coefficient matrix of shape `(..., N, N)`.
+ * @param b - Values of shape `(N,)` or `(..., N, M)`.
+ * @returns Solution `x` of shape `(..., N)` or `(..., N, M)`.
+ */
+declare function solve(a: ArrayLike, b: ArrayLike): Array;
+//#endregion
+//#region src/library/numpy/dtype-info.d.ts
+/** @inline */
+type FInfo = Readonly<{
+  /** The number of bits occupied by the type. */
+  bits: number;
+  /** Returns the _dtype_ for which finfo returns information. */
+  dtype: DType;
+  /** The difference between 1.0 and the next smallest representable float larger than 1.0. */
+  eps: number;
+  /** The difference between 1.0 and the next largest representable float smaller than 1.0. */
+  epsneg: number;
+  /** The exponent that yields `eps`. */
+  machep: number;
+  /** The largest representable finite number. */
+  max: number;
+  /** The smallest positive power of the base (2) that causes overflow. */
+  maxexp: number;
+  /** The smallest representable (most negative) finite number. */
+  min: number;
+  /** The largest negative power of the base (2) without leading zeros in mantissa. */
+  minexp: number;
+  /** The exponent that yields `epsneg`. */
+  negep: number;
+  /** Number of bits in the exponent portion. */
+  nexp: number;
+  /** Number of bits in the mantissa portion. */
+  nmant: number;
+  /** The approximate number of decimal digits to which this kind of float is precise. */
+  precision: number;
+  /** The approximate decimal resolution, i.e., `10 ** -precision`. */
+  resolution: number;
+  /** The smallest positive normal number. */
+  smallestNormal: number;
+  /** The smallest positive subnormal number. */
+  smallestSubnormal: number;
+}>;
+/** Machine limits for floating-point types. */
+declare function finfo(dtype: DType): FInfo;
+/** @inline */
+type IInfo = Readonly<{
+  /** The number of bits occupied by the type. */
+  bits: number;
+  /** Returns the _dtype_ for which iinfo returns information. */
+  dtype: DType;
+  /** The largest representable integer. */
+  max: number;
+  /** The smallest representable integer. */
+  min: number;
+}>;
+/** 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, dot$1 as dot, dstack, e, einsum, equal, eulerGamma, exp, exp2, expandDims, expm1, eye, numpy_fft_d_exports as fft, flip, fliplr, flipud, float16, float32, float64, floor, fmod, frexp, full, fullLike, greater, greaterEqual, hamming, hann, heaviside, hstack, hypot, identity$1 as identity, inf, inner, int32, isfinite, isinf, isnan, isneginf, isposinf, ldexp, less, lessEqual, numpy_linalg_d_exports as linalg, linspace, log, log10, log1p, log2, matmul, matrixTranspose, max, maximum, mean, meshgrid, min, minimum, moveaxis, multiply, nan, 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, sinh, size, sort, sqrt, square, squeeze, stack, std, subtract, sum, 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, 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$1 as 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, 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 };
 }
 declare const float32 = DType.Float32;
 declare const int32 = DType.Int32;
@@ -732,6 +830,16 @@ declare function argmax(a: ArrayLike, axis?: number, opts?: ReduceOpts): Array;
 declare function cumsum(a: ArrayLike, axis?: number): Array;
 /** Reverse the elements in an array along the given axes. */
 declare function flip(x: ArrayLike, axis?: Axis): Array;
+/**
+ * Split an array into multiple sub-arrays along an axis.
+ *
+ * @param a - The input array to split.
+ * @param indicesOrSections - If an integer, it indicates the number of equal
+ * sections to create along the specified axis. If a list of integers, it
+ * specifies the indices at which to split the array.
+ * @param axis - The axis along which to split the array. Default is 0.
+ */
+declare function split$1(a: ArrayLike, indicesOrSections: number | number[], axis?: number): Array[];
 /**
  * Join a sequence of arrays along an existing axis.
  *
@@ -775,6 +883,8 @@ declare function columnStack(xs: ArrayLike[]): Array;
 declare function flipud(x: ArrayLike): Array;
 /** Flip an array horizontally (axis=1). */
 declare function fliplr(x: ArrayLike): Array;
+/** Interchange two axes of an array. */
+declare function swapaxes(a: ArrayLike, axis1: number, axis2: number): Array;
 /** Transpose the last two dimensions of an array. */
 declare function matrixTranspose(a: ArrayLike): Array;
 /** Return a 1-D flattened array containing the elements of the input. */
@@ -860,6 +970,13 @@ declare function sort(a: ArrayLike, axis?: number): Array;
  * The array is sorted along a specified axis (the last by default).
  */
 declare function argsort(a: ArrayLike, axis?: number): Array;
+/**
+ * Take elements from an array along an axis.
+ *
+ * This is equivalent to advanced indexing with integer indices over that
+ * 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. */
 declare function allclose(actual: Parameters<typeof array>[0], expected: Parameters<typeof array>[0], options?: {
   rtol?: number;
@@ -990,6 +1107,17 @@ declare const heaviside: OwnedFunction<(x1: ArrayLike, x2: ArrayLike) => Array>;
 declare function square(x: ArrayLike): Array;
 /** Element-wise tangent function (takes radians). */
 declare function tan(x: ArrayLike): Array;
+/**
+ * @function
+ * Return the normalized sinc function.
+ *
+ * The sinc function is defined as `sin(πx) / (πx)` for `x != 0`, and `1` for `x = 0`.
+ * This is the normalized sinc function commonly used in signal processing.
+ *
+ * **Note:** JVP is not supported at x=0 due to discontinuous derivative. This
+ * requires a custom JVP rule to handle properly (see JAX implementation).
+ */
+declare const sinc: OwnedFunction<(x: ArrayLike) => Array>;
 /** Element-wise inverse cosine function (inverse of cos). */
 declare function acos(x: ArrayLike): Array;
 /**
@@ -1019,6 +1147,20 @@ declare const atan2: OwnedFunction<(y: ArrayLike, x: ArrayLike) => Array>;
 declare function subtract(x: ArrayLike, y: ArrayLike): Array;
 /** Calculates the floating-point division of x by y element-wise. */
 declare function trueDivide(x: ArrayLike, y: ArrayLike): Array;
+/**
+ * Return the largest integer smaller or equal to the division of the inputs.
+ *
+ * The result is always rounded towards negative infinity.
+ *
+ * For floating-point inputs, this is equivalent to `floor(x / y)`.
+ * For integer inputs, we use `(x - remainder(x, y)) / y` to handle
+ * negative values correctly (note: may overflow near int32 boundaries).
+ *
+ * @param x - Dividend array.
+ * @param y - Divisor array.
+ * @returns Element-wise floor division of x by y.
+ */
+declare function floorDivide(x: ArrayLike, y: ArrayLike): Array;
 /**
  * @function
  * Calculate element-wise floating-point modulo operation.
@@ -1029,6 +1171,16 @@ declare const fmod: OwnedFunction<(x: ArrayLike, y: ArrayLike) => Array>;
  * Calculate element-wise remainder of the division (matches sign of y).
  */
 declare const remainder: OwnedFunction<(x: ArrayLike, y: ArrayLike) => Array>;
+/**
+ * Return element-wise quotient and remainder simultaneously.
+ *
+ * Equivalent to `[floorDivide(x, y), remainder(x, y)]`.
+ *
+ * @param x - Dividend array.
+ * @param y - Divisor array.
+ * @returns Tuple of [quotient, remainder].
+ */
+declare function divmod(x: ArrayLike, y: ArrayLike): [Array, Array];
 /** Round input to the nearest integer towards zero. */
 declare function trunc(x: ArrayLike): Array;
 /**
@@ -1139,7 +1291,11 @@ declare function std(x: ArrayLike, axis?: Axis, opts?: {
   correction?: number;
 } & ReduceOpts): Array;
 /** Estimate the sample covariance of a set of variables. */
-declare function cov(x: ArrayLike, y?: ArrayLike): Array;
+declare function cov(x: ArrayLike, y?: ArrayLike | null, {
+  rowvar
+}?: {
+  rowvar?: boolean;
+}): Array;
 /** Compute the Pearson correlation coefficients (in range `[-1, 1]`). */
 declare function corrcoef(x: ArrayLike, y?: ArrayLike): Array;
 /** Test element-wise for positive or negative infinity, return bool array. */
@@ -1353,6 +1509,8 @@ declare enum Primitive {
   PoolTranspose = "pool_transpose",
   Compare = "compare",
   Where = "where",
+  Concatenate = "concatenate",
+  Split = "split",
   RandomBits = "random_bits",
   Gather = "gather",
   Transpose = "transpose",
@@ -1369,6 +1527,8 @@ declare enum Primitive {
   // A is upper triangular, A @ X.T = B.T
   Cholesky = "cholesky",
   // A is positive-definite, A = L @ L^T
+  LU = "lu",
+  // LU decomposition with partial pivoting
   Jit = "jit",
 }
 interface PrimitiveParamsImpl extends Record<Primitive, Record<string, any>> {
@@ -1395,6 +1555,13 @@ interface PrimitiveParamsImpl extends Record<Primitive, Record<string, any>> {
   [Primitive.Compare]: {
     op: CompareOp;
   };
+  [Primitive.Concatenate]: {
+    axis: number;
+  };
+  [Primitive.Split]: {
+    axis: number;
+    sizes: number[];
+  };
   [Primitive.RandomBits]: {
     shape: number[];
     mode: "xor" | 0 | 1;
@@ -1422,14 +1589,14 @@ interface PrimitiveParamsImpl extends Record<Primitive, Record<string, any>> {
   [Primitive.Pad]: {
     width: Pair[];
   };
+  [Primitive.TriangularSolve]: {
+    unitDiagonal: boolean;
+  };
   [Primitive.Jit]: {
     name: string;
     jaxpr: Jaxpr;
     numConsts: number;
   };
-  [Primitive.TriangularSolve]: {
-    unitDiagonal: boolean;
-  };
 }
 /** Type of parameters taken by each primitive. */
 type PrimitiveParams<T extends Primitive> = T extends keyof PrimitiveParamsImpl ? PrimitiveParamsImpl[T] : Record<string, never>;
@@ -1570,6 +1737,7 @@ declare abstract class Tracer {
   neg(): this;
   add(other: this | TracerValue): this;
   mul(other: this | TracerValue): this;
+  mod(other: this | TracerValue): this;
   greater(other: this | TracerValue): this;
   less(other: this | TracerValue): this;
   equal(other: this | TracerValue): this;
@@ -1672,6 +1840,7 @@ declare class ShapedArray implements AbstractValue {
   static fromAval(aval: AbstractValue): ShapedArray;
   get ndim(): number;
   get size(): number;
+  scalar(): ShapedArray;
   toString(): string;
   equals(other: ShapedArray): boolean;
 }
@@ -1739,6 +1908,8 @@ declare class Array extends Tracer {
   toString(): string;
   get device(): Device;
   get ref(): this;
+  /** Get the current reference count (for debugging memory management). */
+  get refCount(): number;
   dispose(): void;
   /**
    * Convert this array into a primitive value.
@@ -1887,8 +2058,43 @@ declare function linspace(start: number, stop: number, num?: number, endpoint?:
   dtype,
   device
 }?: DTypeAndDevice): Array;
+/**
+ * Return numbers spaced evenly on a log scale.
+ *
+ * In linear space, the sequence starts at `base ** start` and ends at
+ * `base ** stop` (see `endpoint` below).
+ *
+ * @param start - `base ** start` is the starting value of the sequence.
+ * @param stop - `base ** stop` is the final value of the sequence, unless `endpoint` is false.
+ * @param num - Number of samples to generate. Default is 50.
+ * @param endpoint - If true, `stop` is the last sample. Otherwise, it is not included. Default is true.
+ * @param base - The base of the log space. Default is 10.
+ * @returns Array of evenly spaced values on a log scale.
+ */
+declare function logspace(start: number, stop: number, num?: number, endpoint?: boolean, base?: number, {
+  dtype,
+  device
+}?: DTypeAndDevice): Array;
+//#endregion
+//#region src/frontend/linearize.d.ts
+/** @inline */
+type GradOpts = {
+  /**
+   * Integer or sequence of integers. Specifies which positional argument(s) to
+   * differentiate with respect to.
+   *
+   * Defaults to `0` (the first argument).
+   */
+  argnums?: number | number[];
+  /**
+   * The input function returns a pair of `[out, aux]` including an auxiliary
+   * value. This `aux` is not differentiated, but is returned alongside the
+   * gradient when evaluating the function.
+   */
+  hasAux?: boolean;
+};
 declare namespace lax_linalg_d_exports {
-  export { cholesky, triangularSolve };
+  export { cholesky, lu, triangularSolve };
 }
 /**
  * Compute the Cholesky decomposition of a symmetric positive-definite matrix.
@@ -1921,6 +2127,32 @@ declare function cholesky(a: ArrayLike, {
 }?: {
   upper?: boolean;
 }): Array;
+/**
+ * LU decomposition with partial pivoting.
+ *
+ * Computes the matrix decomposition: `P @ A = L @ U`, where `P` is a
+ * permutation of the rows of `A`, `L` is lower-triangular with unit diagonal,
+ * and `U` is upper-triangular.
+ *
+ * @param x - A batch of matrices with shape `[..., m, n]`.
+ *
+ * @returns A tuple `(lu, pivots, permutation)` where:
+ * - `lu`: combined lower and upper triangular matrices.
+ * - `pivots`: an array of pivot indices with shape `[..., min(m, n)]`.
+ * - `permutation`: the permutation generated by pivots with shape `[..., m]`.
+ *
+ * @example
+ * ```ts
+ * import { lax, numpy as np } from "@jax-js/jax";
+ *
+ * const A = np.array([[4., 3.], [6., 3.]]);
+ * const [lu, pivots, permutation] = lax.linalg.lu(A);
+ * // lu ≈ [[6., 3.], [0.6666667, 1.0]]
+ * // pivots = [1, 1]
+ * // permutation = [1, 0]
+ * ```
+ */
+declare function lu(x: ArrayLike): [Array, Array, Array];
 /**
  * Solve a triangular linear system.
  *
@@ -1951,7 +2183,7 @@ declare function triangularSolve(a: ArrayLike, b: ArrayLike, {
   unitDiagonal?: boolean;
 }): Array;
 declare namespace lax_d_exports {
-  export { DotDimensionNumbers, PaddingType, conv, convGeneralDilated, convWithGeneralPadding, dot, erf, erfc, lax_linalg_d_exports as linalg, reduceWindow, stopGradient };
+  export { DotDimensionNumbers, PaddingType, conv, convGeneralDilated, convTranspose, convWithGeneralPadding, dot, erf, erfc, lax_linalg_d_exports as linalg, reduceWindow, stopGradient };
 }
 /**
  * Dimension numbers for general `dot()` primitive.
@@ -1989,7 +2221,11 @@ type PaddingType = "VALID" | "SAME" | "SAME_LOWER" | Pair[];
  * The semantics of this operation mimic the `jax.lax.conv_general_dilated`
  * function in JAX, which wraps XLA's general convolution operator.
  *
- * Grouped convolutions are not supported right now.
+ * @param lhs - Input tensor; shape `[N, C_in, ...xs]`
+ * @param rhs - Convolution kernel; shape `[C_out, C_in / G, ...ks]`
+ * @param windowStrides - Strides for each spatial dimension
+ * @param padding - Padding for each spatial dimension, or a string
+ *   (`"VALID"`, `"SAME"`, or `"SAME_LOWER"`)
  */
 declare function convGeneralDilated(lhs: Array, rhs: Array, windowStrides: number[], padding: PaddingType, {
   lhsDilation,
@@ -2004,6 +2240,37 @@ declare function convGeneralDilated(lhs: Array, rhs: Array, windowStrides: numbe
 declare function convWithGeneralPadding(lhs: Array, rhs: Array, windowStrides: number[], padding: PaddingType, lhsDilation?: number[], rhsDilation?: number[]): Array;
 /** Convenience wrapper around `convGeneralDilated`. */
 declare function conv(lhs: Array, rhs: Array, windowStrides: number[], padding: PaddingType): Array;
+/**
+ * Convenience wrapper for calculating the N-d convolution "transpose".
+ *
+ * This function directly calculates a fractionally strided conv rather than
+ * indirectly calculating the gradient (transpose) of a forward convolution.
+ * It is equivalent to the JAX version, except:
+ *
+ * - The `use_consistent_padding` option is not available. We only have the
+ *   consistent padding case (JAX version >0.8.4).
+ * - The order of dimensions matches `lax.conv_general_dilated`.
+ *
+ * Unlike PyTorch/TensorFlow, by default we don't reverse the kernel's spatial
+ * dimensions or the `(C_out, C_in)` axis order. To get this behavior, set
+ * `transposeKernel` to true.
+ *
+ * @param lhs - Input tensor; shape `[N, C_in, ...xs]`
+ * @param rhs - Convolution kernel; shape `[C_out, C_in, ...ks]`
+ * @param strides - Sequence of n integers, sets fractional stride
+ * @param padding - Apply padding of `dilation * (kernel_size - 1) - padding` to
+ *   each side of the input, so it acts like gradient of `conv()`
+ * @param rhsDilation - Atrous dilation for the kernel
+ * @param transposeKernel - Flip spatial axes and swap the input/output channels
+ *   of the kernel; its shape should be `[C_in, C_out, ...ks]`
+ */
+declare function convTranspose(lhs: Array, rhs: Array, strides: number[], padding: PaddingType, {
+  rhsDilation,
+  transposeKernel
+}?: {
+  rhsDilation?: number[];
+  transposeKernel?: boolean;
+}): Array;
 /** Reduce a computation over padded windows. */
 declare function reduceWindow(operand: Array, computation: (x: Array) => Array, windowDimensions: number[], windowStrides?: number[]): Array;
 /** The error function: `erf(x) = 2/sqrt(pi) * int[0..x] exp(-t^2) dt`. */
@@ -2023,7 +2290,7 @@ declare function erfc(x: ArrayLike): Array;
  */
 declare function stopGradient(x: ArrayLike): Array;
 declare namespace nn_d_exports {
-  export { celu, elu, gelu, glu, hardSigmoid, hardSilu, hardSilu as hardSwish, hardTanh, identity, leakyRelu, logSigmoid, logSoftmax, logmeanexp, logsumexp, mish, oneHot, relu, relu6, selu, sigmoid, silu, softSign, softmax, softplus, sparsePlus, sparseSigmoid, squareplus, standardize, silu as swish };
+  export { celu, dotProductAttention, elu, gelu, glu, hardSigmoid, hardSilu, hardSilu as hardSwish, hardTanh, identity, leakyRelu, logSigmoid, logSoftmax, logmeanexp, logsumexp, mish, oneHot, relu, relu6, selu, sigmoid, silu, softSign, softmax, softplus, sparsePlus, sparseSigmoid, squareplus, standardize, silu as swish };
 }
 /**
  * Rectified Linear Unit (ReLU) activation function:
@@ -2220,11 +2487,61 @@ declare function standardize(x: ArrayLike, axis?: Axis, opts?: {
  * ```
  */
 declare function oneHot(x: Array, numClasses: number): Array;
+/**
+ * Scaled dot product attention (SDPA).
+ *
+ * Computes `softmax((Q @ K^T) / sqrt(d) + bias) @ V`, where `Q` is the query,
+ * `K` is the key, `V` is the value, and `d` is the dimensionality of each key
+ * and query vector.
+ *
+ * Multi-query attention is applied when input `key` and `value` tensors have
+ * fewer heads than `query`.
+ *
+ * We use the following uppercase letters to denote array shapes:
+ * - `B` = batch size
+ * - `S` = length of key/value sequences (source)
+ * - `L` = length of query sequences
+ * - `N` = number of attention heads
+ * - `H` = dimensionality of each attention head
+ * - `K` = number of key/value heads (for grouped-query attention)
+ *
+ * The batch size `B` may be omitted, which is equivalent to `B = 1`. In this
+ * case it must be omitted from all inputs.
+ *
+ * @param query - Query array; shape `[B, L, N, H]`
+ * @param key - Key array; shape `[B, S, K, H]`
+ * @param value - Value array; same shape as `key`
+ * @param opts.bias - Optional bias to add to the attention logits; shape
+ *   `[B, N, L, S]` or broadcastable to it.
+ * @param opts.mask - Optional mask to apply to the attention logits; should be
+ *   a boolean array broadcastable to `[B, N, L, S]`, where `true` indicates
+ *   the element should take part in attention.
+ * @param opts.scale - Scaling factor override, default is `1 / sqrt(H)`.
+ * @param opts.isCausal - If true, applies a casual mask.
+ * @param opts.querySeqLengths - Optional sequence lengths for the queries;
+ *   shape `(B,)`. Taken from the beginning of the tensor.
+ * @param opts.keyValueSeqLengths - Optional sequence lengths for the keys and
+ *   values; shape `(B,)`. Taken from the beginning of the tensor.
+ * @param opts.localWindowSize - If specified, applies a local attention window
+ *   of the given size. Can be a single number or a tuple `[left, right]`.
+ *
+ * @returns The result of the attention operation; shape is the same as query
+ *   `[B, L, N, H]`, or `[L, N, H]` if `B` is omitted.
+ */
+declare function dotProductAttention(query: ArrayLike, key: ArrayLike, value: ArrayLike, opts?: {
+  bias?: ArrayLike;
+  mask?: ArrayLike;
+  scale?: number;
+  isCausal?: boolean;
+  querySeqLengths?: ArrayLike;
+  keyValueSeqLengths?: ArrayLike;
+  localWindowSize?: number | [number, number];
+}): Array;
 declare namespace random_d_exports {
-  export { bernoulli, bits, cauchy, exponential, gumbel, key, laplace, normal, split, uniform };
+  export { bernoulli, bits, cauchy, exponential, gumbel, key, laplace, multivariateNormal, normal, split, uniform };
 }
 /** Create a pseudo-random number generator (PRNG) key from 32-bit integer seed. */
-declare function key(seed: number): Array;
+declare function key(seed: ArrayLike): Array;
 /** Splits a PRNG key into `num` new keys by adding a leading axis. */
 declare function split(key: Array, num?: number | number[]): Array;
 /** Sample uniform bits in the form of unsigned integers. */
@@ -2271,6 +2588,23 @@ declare const gumbel: OwnedFunction<(key: ArrayLike, shape?: number[] | undefine
  * Inverting: `x = -sign(u - 0.5) * log(1 - 2 * |u - 0.5|)`.
  */
 declare const laplace: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined) => Array>;
+/**
+ * @function
+ * Sample multivariate normal random values with given mean and covariance.
+ *
+ * The values are returned with the given shape, along with the final dimension
+ * used to represent the n-dimensional multivariate normal factors.
+ *
+ * This uses Cholesky decomposition on the covariance matrix.
+ *
+ * - `key` - PRNG key
+ * - `mean` - Mean vector of shape `[..., n]`
+ * - `cov` - Covariance of shape `[..., n, n]`, must be positive-definite
+ * - `shape` - Result batch shape, must be broadcastable with
+ *            `mean.shape[:-1]` and `cov.shape[:-2]`
+ * @returns Random samples of shape `[...shape, n]`
+ */
+declare const multivariateNormal: OwnedFunction<(key: ArrayLike, mean: ArrayLike, cov: ArrayLike, shape?: number[] | undefined) => Array>;
 /**
  * @function
  * Sample random values according to `p(x) = 1/sqrt(2pi) * exp(-x^2/2)`.
@@ -2294,7 +2628,9 @@ declare const logit: OwnedFunction<(x: ArrayLike) => Array>;
  * @function
  * Compute the forward-mode Jacobian-vector product for a function.
  */
-declare const jvp: <F extends (...args: any[]) => JsTree<Array>>(f: F, primals: MapJsTree<Parameters<F>, Array, ArrayLike>, tangents: MapJsTree<Parameters<F>, Array, ArrayLike>) => [ReturnType<F>, ReturnType<F>];
+declare const jvp: <F extends (...args: any[]) => JsTree<Array>, HA extends boolean = false>(f: F, primals: MapJsTree<Parameters<F>, Array, ArrayLike>, tangents: MapJsTree<Parameters<F>, Array, ArrayLike>, opts?: {
+  hasAux?: HA;
+}) => HA extends true ? ReturnType<F> extends [infer Out, infer Aux] ? [Out, Out, Aux] : never : [ReturnType<F>, ReturnType<F>];
 /**
  * @function
  * Vectorize an operation on a batched axis for one or more inputs.
@@ -2336,28 +2672,100 @@ declare const jit: <F extends (...args: any[]) => JsTree<Array>>(f: F, opts?: Ji
  * Produce a local linear approximation to a function at a point using jvp() and
  * partial evaluation.
  */
-declare const linearize: <F extends (...args: any[]) => JsTree<Array>>(f: F, ...primals: MapJsTree<Parameters<F>, Array, ArrayLike>) => [ReturnType<F>, (...tangents: MapJsTree<Parameters<F>, Array, ArrayLike>) => ReturnType<F>];
+declare const linearize: <F extends (...args: any[]) => JsTree<Array>, HA extends boolean = false>(f: F, primals: MapJsTree<Parameters<F>, Array, ArrayLike>, opts?: {
+  hasAux?: HA;
+}) => HA extends true ? ReturnType<F> extends [infer Out, infer Aux] ? [Out, OwnedFunction<(...tangents: MapJsTree<Parameters<F>, Array, ArrayLike>) => Out>, Aux] : never : [ReturnType<F>, OwnedFunction<(...tangents: MapJsTree<Parameters<F>, Array, ArrayLike>) => ReturnType<F>>];
 /**
  * @function
  * Calculate the reverse-mode vector-Jacobian product for a function.
+ *
+ * The return value is a tuple of `[out, vjpFn]`, where `out` is the output of
+ * `f(primals)`, and `vjpFn` is a function that takes in cotangents for each
+ * output and returns the cotangents for each input.
+ *
+ * When `{ hasAux: true }` is passed, the function `f` is expected to return an
+ * `[out, aux]` tuple, and `vjp` returns `[out, vjpFn, aux]`.
+ *
+ * @example
+ * ```ts
+ * const [y, vjpFn] = vjp(f, [x]);
+ *
+ * // With hasAux
+ * const [y, vjpFn, aux] = vjp(f, [x], { hasAux: true });
+ * ```
  */
-declare const vjp: <F extends (...args: any[]) => JsTree<Array>>(f: F, ...primals: MapJsTree<Parameters<F>, Array, ArrayLike>) => [ReturnType<F>, (cotangents: MapJsTree<ReturnType<F>, Array, ArrayLike>) => MapJsTree<Parameters<F>, ArrayLike, Array>];
+declare const vjp: <F extends (...args: any[]) => JsTree<Array>, const HA extends boolean = false>(f: F, primals: MapJsTree<Parameters<F>, Array, ArrayLike>, opts?: {
+  hasAux?: HA;
+}) => HA extends true ? ReturnType<F> extends [infer Out, infer Aux] ? [Out, OwnedFunction<(cotangents: MapJsTree<Out, Array, ArrayLike>) => MapJsTree<Parameters<F>, ArrayLike, Array>>, Aux] : never : [ReturnType<F>, OwnedFunction<(cotangents: MapJsTree<ReturnType<F>, Array, ArrayLike>) => MapJsTree<Parameters<F>, ArrayLike, Array>>];
+/** @inline */
+type GradOutputType<I, F extends (...args: any[]) => any> = MapJsTree<I extends undefined ? Parameters<F>[0] : I extends number ? Parameters<F>[I] : I extends number[] ? { [K in keyof I]: I[K] extends number ? Parameters<F>[I[K]] : never } : never, ArrayLike, Array>;
 /**
  * @function
  * Compute the gradient of a scalar-valued function `f` with respect to its
  * first argument.
+ *
+ * Pass in different `argnums` to differentiate with respect to other
+ * arguments. If a tuple is provided, the return value will be a tuple of
+ * gradients corresponding to each argument index.
+ *
+ * When `{ hasAux: true }` is passed, the function `f` is expected to return a
+ * `[out, aux]` tuple, and the return value will be `[gradient, aux]`.
+ *
+ * @example
+ * ```ts
+ * const gradient = grad(f)(x);
+ *
+ * // With `argnums`
+ * const [gradientX, gradientZ] = grad(f, { argnums: [0, 2] })(x, y, z);
+ *
+ * // With `hasAux`
+ * const [gradient, aux] = grad(f, { hasAux: true })(x);
+ * ```
  */
-declare const grad: <F extends (...args: any[]) => JsTree<Array>>(f: F) => (...primals: MapJsTree<Parameters<F>, Array, ArrayLike>) => MapJsTree<Parameters<F>[0], ArrayLike, Array>;
+declare const grad: <F extends (...args: any[]) => JsTree<Array>, const I extends undefined | number | number[] = undefined, const HA extends boolean = false>(f: F, opts?: Omit<GradOpts, "argnums" | "hasAux"> & {
+  argnums?: I;
+  hasAux?: HA;
+}) => (...primals: MapJsTree<Parameters<F>, Array, ArrayLike>) => HA extends true ? ReturnType<F> extends [any, infer Aux] ? [GradOutputType<I, F>, Aux] : never : GradOutputType<I, F>;
 /**
  * @function
  * Create a function that evaluates both `f` and the gradient of `f`.
+ *
+ * When `{ hasAux: true }` is passed, the function `f` is expected to return an
+ * `[out, aux]` tuple, and the return value will be `[[out, aux], gradient]`.
+ *
+ * @example
+ * ```ts
+ * // Without hasAux
+ * const [value, gradient] = valueAndGrad(f)(x);
+ *
+ * // With hasAux
+ * const [[value, aux], gradient] = valueAndGrad(f, { hasAux: true })(x);
+ * ```
  */
-declare const valueAndGrad: <F extends (...args: any[]) => JsTree<Array>>(f: F) => (...primals: MapJsTree<Parameters<F>, Array, ArrayLike>) => [ReturnType<F>, MapJsTree<Parameters<F>[0], ArrayLike, Array>];
+declare const valueAndGrad: <F extends (...args: any[]) => JsTree<Array>, const I extends undefined | number | number[] = undefined, const HA extends boolean = false>(f: F, opts?: Omit<GradOpts, "argnums"> & {
+  argnums?: I;
+  hasAux?: HA;
+}) => (...primals: MapJsTree<Parameters<F>, Array, ArrayLike>) => [ReturnType<F>, GradOutputType<I, F>];
 /**
  * @function
  * Compute the Jacobian evaluated row-by-row by reverse-mode AD.
  */
 declare const jacrev: typeof jacfwd;
+/**
+ * @function
+ * Compute the Hessian matrix of a scalar-valued function.
+ *
+ * The Hessian is the matrix of second-order partial derivatives of a function.
+ * This is implemented as `jacfwd(grad(f))`.
+ *
+ * @example
+ * ```ts
+ * const f = (x: np.Array) => np.sum(x.ref.mul(x.ref).mul(x)); // x^3
+ * const H = hessian(f)(np.array([1, 2, 3]));
+ * // H[i,j] = d^2f / dx_i dx_j
+ * ```
+ */
+declare const hessian: <F extends (x: Array) => Array>(f: F) => (...args: MapJsTree<Parameters<F>, Array, ArrayLike>) => ReturnType<F>;
 /**
  * Wait until all `Array` leaves are ready by calling `Array.blockUntilReady()`.
  *
@@ -2380,4 +2788,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, 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, random_d_exports as random, scipy_special_d_exports as scipySpecial, setDebug, tree_d_exports as tree, valueAndGrad, vjp, vmap };