@jax-js/jax 0.1.4 → 0.1.5

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, 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.
  *
@@ -860,6 +968,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 +1105,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 +1145,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 +1169,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 +1289,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 +1507,8 @@ declare enum Primitive {
   PoolTranspose = "pool_transpose",
   Compare = "compare",
   Where = "where",
+  Concatenate = "concatenate",
+  Split = "split",
   RandomBits = "random_bits",
   Gather = "gather",
   Transpose = "transpose",
@@ -1369,6 +1525,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 +1553,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;
@@ -1570,6 +1735,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 +1838,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 +1906,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 +2056,25 @@ 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;
 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 +2107,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.
  *
@@ -2221,10 +2433,10 @@ declare function standardize(x: ArrayLike, axis?: Axis, opts?: {
  */
 declare function oneHot(x: Array, numClasses: 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 +2483,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)`.