@jax-js/jax 0.1.3 → 0.1.5

package/dist/index.d.ts

@@ -121,7 +121,6 @@ declare class ShapeTracker {
   /** Like pad(), but allows for negative values. */
   padOrShrink(arg: Pair[]): ShapeTracker;
 }
-//# sourceMappingURL=shape.d.ts.map
 //#endregion
 //#region src/utils.d.ts
 /**
@@ -404,8 +403,81 @@ declare class Reduction implements FpHashable {
 }
 /** Expression for accessing `indices` in input array with the given shape. */
 //#endregion
+//#region src/routine.d.ts
+/**
+ * Advanced operations that don't fit into the `AluExp` compiler representation.
+ *
+ * Some routines like iterative matrix algorithms, FFTs, or sorting may not be
+ * easy to express efficiently as a `Kernel` object. These also tend to be
+ * somewhat expensive, so the benefit of kernel fusion and inlining is less
+ * relevant.
+ *
+ * For these operations, we dispatch them as a custom operation on the backend,
+ * which each backend implements in a specific way. These are listed in the
+ * `Routines` enum below.
+ *
+ * Routines cannot be fused into other kernels and always operate on contiguous
+ * arrays (default `ShapeTracker`).
+ */
+declare class Routine {
+  /** The name of the routine. */
+  readonly name: Routines;
+  /** Dtype and shape of the inputs and outputs. */
+  readonly type: RoutineType;
+  /** Extra parameters specific to the routine. */
+  readonly params?: any | undefined;
+  constructor(/** The name of the routine. */
+  name: Routines, /** Dtype and shape of the inputs and outputs. */
+  type: RoutineType, /** Extra parameters specific to the routine. */
+  params?: any | undefined);
+}
+/** One of the valid `Routine` that can be dispatched to backend. */
+declare enum Routines {
+  /** Stable sorting algorithm along the last axis. */
+  Sort = "Sort",
+  /** Returns `int32` indices of the stably sorted array. */
+  Argsort = "Argsort",
+  /**
+   * 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.
+   *
+   * 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[][];
+  inputDtypes: DType[];
+  outputShapes: number[][];
+  outputDtypes: DType[];
+}
+//#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;
@@ -441,9 +513,13 @@ interface Backend {
   /** Read a range of bytes from a buffer, blocking variant. */
   readSync(slot: Slot, start?: number, count?: number): Uint8Array<ArrayBuffer>;
   /** Prepare an expression to be executed later. */
-  prepare(kernel: Kernel): Promise<Executable>;
+  prepareKernel(kernel: Kernel): Promise<Executable>;
   /** Prepare an expression to be executed later, blocking variant. */
-  prepareSync(kernel: Kernel): Executable;
+  prepareKernelSync(kernel: Kernel): Executable;
+  /** Prepare an advanced routine to be executed later. */
+  prepareRoutine(routine: Routine): Promise<Executable>;
+  /** Prepare an advanced routine to be executed later, blocking variant. */
+  prepareRoutineSync(routine: Routine): Executable;
   /**
    * Run a backend operation that was previously prepared.
    *
@@ -454,14 +530,140 @@ interface Backend {
   dispatch(exe: Executable, inputs: Slot[], outputs: Slot[]): void;
 }
 declare class Executable<T = any> {
-  readonly kernel: Kernel;
-  /** Extra data specific to the backend running this kernel. */
+  /** The `Kernel` or `Routine` that was prepared. */
+  readonly source: Kernel | Routine;
+  /** Extra data specific to the backend running this executable. */
   readonly data: T;
-  constructor(kernel: Kernel, /** Extra data specific to the backend running this kernel. */
+  constructor(/** The `Kernel` or `Routine` that was prepared. */
+  source: Kernel | Routine, /** Extra data specific to the backend running this executable. */
   data: T);
 }
+declare namespace numpy_fft_d_exports {
+  export { ComplexPair, fft, ifft };
+}
+/**
+ * A pair of arrays representing real and imaginary part `a + bj`. Both arrays
+ * must have the same shape.
+ */
+type ComplexPair = {
+  real: Array;
+  imag: Array;
+};
+/**
+ * Compute a one-dimensional discrete Fourier transform.
+ *
+ * Currently, the size of the axis must be a power of two.
+ */
+declare function fft(a: ComplexPair, axis?: number): ComplexPair;
+/**
+ * Compute a one-dimensional inverse discrete Fourier transform.
+ *
+ * Currently, the size of the axis must be a power of two.
+ */
+declare function ifft(a: ComplexPair, axis?: number): ComplexPair;
+declare namespace numpy_linalg_d_exports {
+  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.
+ *
+ * This is like `jax.lax.linalg.cholesky()`, except with an option to symmetrize
+ * the input matrix, which is on by default.
+ */
+declare function cholesky$1(a: ArrayLike, {
+  upper,
+  symmetrizeInput
+}?: {
+  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.
+ *
+ * For overdetermined systems, this finds the `x` that minimizes `norm(ax - b)`.
+ * For underdetermined systems, this finds the minimum-norm solution for `x`.
+ *
+ * This currently uses Cholesky decomposition to solve the normal equations,
+ * under the hood. The method is not as robust as QR or SVD.
+ *
+ * @param a coefficient matrix of shape `(M, N)`
+ * @param b right-hand side of shape `(M,)` or `(M, K)`
+ * @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, abs, absolute, acos, acosh, add, allclose, arange, arccos, arccosh, arcsinh, arctan, arctan2, arctanh, argmax, argmin, array, asin, asinh, astype, atan, atan2, atanh, bool, broadcastArrays, broadcastShapes, broadcastTo, cbrt, ceil, clip, columnStack, concatenate, cos, cosh, cumsum, cumulativeSum, deg2rad, degrees, diag, diagonal, divide, dot$1 as dot, dstack, e, einsum, equal, eulerGamma, exp, exp2, expm1, eye, 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, linspace, log, log10, log1p, log2, matmul, max, maximum, mean, meshgrid, min, minimum, moveaxis, multiply, nan, ndim, negative, notEqual, ones, onesLike, outer, pad, permuteDims, pi, positive, pow, power, prod, promoteTypes, ptp, rad2deg, radians, ravel, reciprocal, remainder, repeat, reshape, shape$1 as shape, sign, sin, sinh, size, 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;
@@ -587,6 +789,20 @@ declare function prod(a: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
 declare function min(a: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
 /** Return the maximum of array elements along a given axis. */
 declare function max(a: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
+/**
+ * Test whether all array elements along a given axis evaluate to True.
+ *
+ * Returns a boolean array with the same shape as `a` with the specified axis
+ * removed. If axis is None, returns a scalar.
+ */
+declare function all(a: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
+/**
+ * Test whether any array element along a given axis evaluates to True.
+ *
+ * Returns a boolean array with the same shape as `a` with the specified axis
+ * removed. If axis is None, returns a scalar.
+ */
+declare function any(a: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
 /** Return the peak-to-peak range along a given axis (`max - min`). */
 declare function ptp(a: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
 /** Compute the average of the array elements along the specified axis. */
@@ -612,10 +828,18 @@ declare function argmax(a: ArrayLike, axis?: number, opts?: ReduceOpts): Array;
  * two-phase parallel reduction algorithm.
  */
 declare function cumsum(a: ArrayLike, axis?: number): Array;
-/** @function Alternative name for `jax.numpy.cumsum()`. */
-declare const cumulativeSum: typeof cumsum;
 /** 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.
  *
@@ -659,12 +883,29 @@ declare function columnStack(xs: ArrayLike[]): Array;
 declare function flipud(x: ArrayLike): Array;
 /** Flip an array horizontally (axis=1). */
 declare function fliplr(x: ArrayLike): Array;
-/** @function Alternative name for `numpy.transpose()`. */
-declare const permuteDims: (x: ArrayLike, perm?: 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. */
 declare function ravel(a: ArrayLike): Array;
 /** Remove one or more length-1 axes from an array. */
 declare function squeeze(a: ArrayLike, axis?: Axis): Array;
+/**
+ * Expand the shape of an array by inserting new axes of length 1.
+ *
+ * @param a - Input array.
+ * @param axis - Position(s) in the expanded axes where the new axis (or axes)
+ *   is placed. Can be a single integer or an array of integers.
+ * @returns Array with the number of dimensions increased.
+ *
+ * @example
+ * ```ts
+ * const x = np.array([1, 2]);
+ * np.expandDims(x, 0); // Shape [1, 2]
+ * np.expandDims(x, 1); // Shape [2, 1]
+ * np.expandDims(x, [0, 2]); // Shape [1, 2, 1]
+ * ```
+ */
+declare function expandDims(a: ArrayLike, axis: number | number[]): Array;
 /**
  * Repeat each element of an array after themselves.
  *
@@ -711,6 +952,29 @@ declare function diagonal(a: ArrayLike, offset?: number, axis1?: number, axis2?:
 declare function diag(v: ArrayLike, k?: number): Array;
 /** Calculate the sum of the diagonal of an array along the given axes. */
 declare function trace(a: ArrayLike, offset?: number, axis1?: number, axis2?: number): Array;
+/**
+ * Return a sorted copy of an array.
+ *
+ * The array is sorted along a specified axis (the last by default). This may be
+ * an unstable sort, and it dispatches to device-specific implementation.
+ */
+declare function sort(a: ArrayLike, axis?: number): Array;
+/**
+ * Return indices that would sort an array. This may be an unstable sorting
+ * algorithm; it need not preserve order of indices in ties.
+ *
+ * Returns an array of `int32` indices.
+ *
+ * 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;
@@ -782,6 +1046,10 @@ declare function vecdot(x: ArrayLike, y: ArrayLike, {
  * Like vecdot() but flattens the arguments first into vectors.
  */
 declare function vdot(x: ArrayLike, y: ArrayLike): Array;
+/** Convolution of two one-dimensional arrays. */
+declare function convolve(x: Array, y: Array, mode?: "full" | "same" | "valid"): Array;
+/** Correlation of two one dimensional arrays. */
+declare function correlate(x: Array, y: Array, mode?: "full" | "same" | "valid"): Array;
 /**
  * Return a tuple of coordinate matrices from coordinate vectors.
  *
@@ -793,21 +1061,6 @@ declare function meshgrid(xs: Array[], {
 }?: {
   indexing?: "xy" | "ij";
 }): Array[];
-/**
- * Return an array with ones on and below the diagonal and zeros elsewhere.
- *
- * If `k` is provided, it specifies the sub-diagonal on and below which the
- * array is filled with ones. `k=0` is the main diagonal, `k<0` is below it, and
- * `k>0` is above it.
- */
-declare function tri(n: number, m?: number, k?: number, {
-  dtype,
-  device
-}?: DTypeAndDevice): Array;
-/** Return the lower triangle of an array. Must be of dimension >= 2. */
-declare function tril(a: ArrayLike, k?: number): Array;
-/** Return the upper triangle of an array. Must be of dimension >= 2. */
-declare function triu(a: ArrayLike, k?: number): Array;
 /**
  * Clip (limit) the values in an array.
  *
@@ -824,8 +1077,6 @@ declare function clip(a: ArrayLike, min?: ArrayLike, max?: ArrayLike): Array;
  * This is the same function as `jax.numpy.abs()`.
  */
 declare function absolute(x: ArrayLike): Array;
-/** @function Alias of `jax.numpy.absolute()`. */
-declare const abs: typeof absolute;
 /** 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). */
@@ -854,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;
 /**
@@ -879,16 +1141,24 @@ declare const hypot: OwnedFunction<(x1: ArrayLike, x2: ArrayLike) => Array>;
  * The output is ill-defined when both x and y are zero.
  */
 declare const atan2: OwnedFunction<(y: ArrayLike, x: ArrayLike) => Array>;
-/** @function Alias of `jax.numpy.acos()`. */
-declare const arccos: typeof acos;
-/** @function Alias of `jax.numpy.atan()`. */
-declare const arctan: (x: ArrayLike) => Array;
-/** @function Alias of `jax.numpy.atan2()`. */
-declare const arctan2: OwnedFunction<(y: ArrayLike, x: ArrayLike) => Array>;
 /** Element-wise subtraction, with broadcasting. */
 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.
@@ -899,8 +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>;
-/** @function Alias of `jax.numpy.trueDivide()`. */
-declare const divide: typeof trueDivide;
+/**
+ * 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;
 /**
@@ -940,8 +1218,6 @@ declare const degrees: typeof rad2deg;
  * Computes first array raised to power of second array, element-wise.
  */
 declare const power: OwnedFunction<(x1: ArrayLike, x2: ArrayLike) => Array>;
-/** @function Alias of `jax.numpy.power()`. */
-declare const pow: OwnedFunction<(x1: ArrayLike, x2: ArrayLike) => Array>;
 /** @function Calculate the element-wise cube root of the input array. */
 declare const cbrt: OwnedFunction<(x: ArrayLike) => Array>;
 /**
@@ -986,12 +1262,6 @@ declare const arccosh: OwnedFunction<(x: ArrayLike) => Array>;
  * `arctanh(x) = 0.5 * ln((1 + x) / (1 - x))`
  */
 declare const arctanh: OwnedFunction<(x: ArrayLike) => Array>;
-/** @function Alias of `jax.numpy.arcsinh()`. */
-declare const asinh: OwnedFunction<(x: ArrayLike) => Array>;
-/** @function Alias of `jax.numpy.arccosh()`. */
-declare const acosh: OwnedFunction<(x: ArrayLike) => Array>;
-/** @function Alias of `jax.numpy.arctanh()`. */
-declare const atanh: OwnedFunction<(x: ArrayLike) => Array>;
 /**
  * Compute the variance of an array.
  *
@@ -1018,6 +1288,14 @@ declare function std(x: ArrayLike, axis?: Axis, opts?: {
   mean?: ArrayLike;
   correction?: number;
 } & ReduceOpts): Array;
+/** Estimate the sample covariance of a set of variables. */
+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. */
 declare function isinf(x: ArrayLike): Array;
 /** Test element-wise for NaN (Not a Number). */
@@ -1031,7 +1309,6 @@ declare function isposinf(x: ArrayLike): Array;
  * Test element-wise for finite values (not infinity or NaN).
  */
 declare const isfinite: OwnedFunction<(x: ArrayLike) => Array>;
-//# sourceMappingURL=numpy.d.ts.map
 declare namespace tree_d_exports {
   export { JsTree, JsTreeDef, MapJsTree, NodeType, dispose, flatten, leaves, map, ref, structure, unflatten };
 }
@@ -1084,7 +1361,7 @@ declare function dispose<Tree extends JsTree<Array>>(tree: Tree | null | undefin
 interface ConvParams {
   vmapDims: number;
   strides: number[];
-  padding: [number, number][];
+  padding: Pair[];
   lhsDilation: number[];
   rhsDilation: number[];
 }
@@ -1165,9 +1442,21 @@ declare class Jaxpr implements FpHashable {
    *  - Remove no-op movement operations.
    */
   simplify(): Jaxpr;
-  /** Flattens nested JitCall in a Jaxpr. Useful for handling jit-of-jit. */
+  /** Flattens nested Jit in a Jaxpr. Useful for handling jit-of-jit. */
   flatten(): Jaxpr;
 }
+/** Jaxpr with a collection of associated, traced constants. */
+declare class ClosedJaxpr {
+  readonly jaxpr: Jaxpr;
+  readonly consts: Tracer[];
+  constructor(jaxpr: Jaxpr, consts: Tracer[]);
+  /** String representation of this Jaxpr. */
+  toString(): string;
+  /** Apply a function to the underlying Jaxpr. */
+  mapJaxpr(f: (jaxpr: Jaxpr) => Jaxpr): ClosedJaxpr;
+  /** Dispose of the constants in this Jaxpr. */
+  dispose(): void;
+}
 /** @inline */
 type JitOpts = {
   staticArgnums?: number[];
@@ -1190,7 +1479,9 @@ declare enum Primitive {
   Mul = "mul",
   Idiv = "idiv",
   Mod = "mod",
-  // uses sign of dividend, C-style, matches JS but not Python
+  // uses sign of numerator, C-style, matches JS but not Python
+  Min = "min",
+  Max = "max",
   Neg = "neg",
   Reciprocal = "reciprocal",
   Floor = "floor",
@@ -1198,7 +1489,6 @@ declare enum Primitive {
   StopGradient = "stop_gradient",
   Cast = "cast",
   Bitcast = "bitcast",
-  RandomBits = "random_bits",
   Sin = "sin",
   Cos = "cos",
   Asin = "asin",
@@ -1208,8 +1498,6 @@ declare enum Primitive {
   Erf = "erf",
   Erfc = "erfc",
   Sqrt = "sqrt",
-  Min = "min",
-  Max = "max",
   Reduce = "reduce",
   Dot = "dot",
   // sum(x*y, axis=-1)
@@ -1219,14 +1507,27 @@ declare enum Primitive {
   PoolTranspose = "pool_transpose",
   Compare = "compare",
   Where = "where",
+  Concatenate = "concatenate",
+  Split = "split",
+  RandomBits = "random_bits",
+  Gather = "gather",
   Transpose = "transpose",
   Broadcast = "broadcast",
   Reshape = "reshape",
   Flip = "flip",
   Shrink = "shrink",
   Pad = "pad",
-  Gather = "gather",
-  JitCall = "jit_call",
+  Sort = "sort",
+  // sort(x, axis=-1)
+  Argsort = "argsort",
+  // argsort(x, axis=-1)
+  TriangularSolve = "triangular_solve",
+  // 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>> {
   [Primitive.Cast]: {
@@ -1252,6 +1553,21 @@ 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;
+  };
+  [Primitive.Gather]: {
+    axis: number[];
+    outDim: number;
+  };
   [Primitive.Transpose]: {
     perm: number[];
   };
@@ -1259,10 +1575,6 @@ interface PrimitiveParamsImpl extends Record<Primitive, Record<string, any>> {
     shape: number[];
     axis: number[];
   };
-  [Primitive.RandomBits]: {
-    shape: number[];
-    mode: "xor" | 0 | 1;
-  };
   [Primitive.Reshape]: {
     shape: number[];
   };
@@ -1275,15 +1587,14 @@ interface PrimitiveParamsImpl extends Record<Primitive, Record<string, any>> {
   [Primitive.Pad]: {
     width: Pair[];
   };
-  [Primitive.Gather]: {
-    axis: number[];
-    outDim: number;
-  };
-  [Primitive.JitCall]: {
+  [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>;
@@ -1424,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;
@@ -1451,7 +1763,7 @@ declare abstract class Tracer {
   sub(other: this | TracerValue): this;
   /** Divide an array by this one. */
   div(other: this | TracerValue): this;
-  /** Return specified diagonals. See `numpy.diagonal` for full docs. */
+  /** Return specified diagonals. See `jax.numpy.diagonal` for full docs. */
   diagonal(offset?: number, axis1?: number, axis2?: number): this;
   /** Flatten the array without changing its data. */
   flatten(): this;
@@ -1470,6 +1782,19 @@ declare abstract class Tracer {
    * ```
    */
   [Symbol.iterator](): IterableIterator<this>;
+  /**
+   * Return a sorted copy of an array in ascending order.
+   *
+   * See `jax.numpy.sort` for full docs.
+   */
+  sort(axis?: number): this;
+  /**
+   * Return the indices that would sort an array. This may not be a stable
+   * sorting algorithm; it need not preserve order of indices in ties.
+   *
+   * See `jax.numpy.argsort` for full docs.
+   */
+  argsort(axis?: number): this;
   /**
    * Slice an array along one or more axes.
    *
@@ -1512,6 +1837,8 @@ declare class ShapedArray implements AbstractValue {
   constructor(shape: number[], dtype: DType, weakType: boolean);
   static fromAval(aval: AbstractValue): ShapedArray;
   get ndim(): number;
+  get size(): number;
+  scalar(): ShapedArray;
   toString(): string;
   equals(other: ShapedArray): boolean;
 }
@@ -1529,12 +1856,12 @@ type ArrayLike = Array | number | boolean;
 declare class PendingExecute {
   #private;
   readonly backend: Backend;
-  readonly kernel: Kernel;
+  readonly source: Kernel | Routine;
   readonly inputs: Slot[];
   readonly outputs: Slot[];
   prepared: Executable | null;
   submitted: boolean;
-  constructor(backend: Backend, kernel: Kernel, inputs: Slot[], outputs: Slot[]);
+  constructor(backend: Backend, source: Kernel | Routine, inputs: Slot[], outputs: Slot[]);
   updateRc(delta: number): void;
   prepare(): Promise<void>;
   prepareSync(): void;
@@ -1566,7 +1893,6 @@ type ArrayConstructorArgs = {
  */
 declare class Array extends Tracer {
   #private;
-  id: number;
   /**
    * @ignore
    * Constructs an array from source, shape and backend. Note that if the source
@@ -1580,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.
@@ -1700,6 +2028,21 @@ declare function arange(start: number, stop?: number, step?: number, {
   dtype,
   device
 }?: DTypeAndDevice): Array;
+/**
+ * Return an array with ones on and below the diagonal and zeros elsewhere.
+ *
+ * If `k` is provided, it specifies the sub-diagonal on and below which the
+ * array is filled with ones. `k=0` is the main diagonal, `k<0` is below it, and
+ * `k>0` is above it.
+ */
+declare function tri(n: number, m?: number, k?: number, {
+  dtype,
+  device
+}?: DTypeAndDevice): Array;
+/** Return the lower triangle of an array. Must be of dimension >= 2. */
+declare function tril(a: ArrayLike, k?: number): Array;
+/** Return the upper triangle of an array. Must be of dimension >= 2. */
+declare function triu(a: ArrayLike, k?: number): Array;
 /**
  * Return evenly spaced numbers over a specified interval.
  *
@@ -1713,8 +2056,114 @@ 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, lu, triangularSolve };
+}
+/**
+ * Compute the Cholesky decomposition of a symmetric positive-definite matrix.
+ *
+ * The Cholesky decomposition of a matrix `A` is:
+ *
+ * - A = L @ L^T  (for upper=false, default)
+ * - A = U^T @ U  (for upper=true)
+ *
+ * where `L` is a lower-triangular matrix and `U` is an upper-triangular matrix.
+ * The input matrix must be symmetric and positive-definite.
+ *
+ * @example
+ * ```ts
+ * import { lax, numpy as np } from "@jax-js/jax";
+ *
+ * const x = np.array([[2., 1.], [1., 2.]]);
+ *
+ * // Lower Cholesky factorization (default):
+ * const L = lax.linalg.cholesky(x);
+ * // L ≈ [[1.4142135, 0], [0.70710677, 1.2247449]]
+ *
+ * // Upper Cholesky factorization:
+ * const U = lax.linalg.cholesky(x, { upper: true });
+ * // U ≈ [[1.4142135, 0.70710677], [0, 1.2247449]]
+ * ```
+ */
+declare function cholesky(a: ArrayLike, {
+  upper
+}?: {
+  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.
+ *
+ * Solves `a @ x = b` (if leftSide=true) or `x @ a = b` (if leftSide=false)
+ * where `a` is a triangular matrix.
+ *
+ * @example
+ * ```ts
+ * import { lax, numpy as np } from "@jax-js/jax";
+ *
+ * const L = np.array([[2., 0.], [1., 3.]]);
+ * const b = np.array([4., 7.]).reshape([2, 1]);
+ *
+ * // Solve L @ x = b
+ * const x = lax.linalg.triangularSolve(L, b, { leftSide: true, lower: true });
+ * // x = [[2.], [5./3.]]
+ * ```
+ */
+declare function triangularSolve(a: ArrayLike, b: ArrayLike, {
+  leftSide,
+  lower,
+  transposeA,
+  unitDiagonal
+}?: {
+  leftSide?: boolean;
+  lower?: boolean;
+  transposeA?: boolean;
+  unitDiagonal?: boolean;
+}): Array;
 declare namespace lax_d_exports {
-  export { DotDimensionNumbers, PaddingType, conv, convGeneralDilated, convWithGeneralPadding, dot, erf, erfc, reduceWindow, stopGradient };
+  export { DotDimensionNumbers, PaddingType, conv, convGeneralDilated, convWithGeneralPadding, dot, erf, erfc, lax_linalg_d_exports as linalg, reduceWindow, stopGradient };
 }
 /**
  * Dimension numbers for general `dot()` primitive.
@@ -1745,7 +2194,7 @@ declare function dot(lhs: Array, rhs: Array, {
   lhsBatchDims: lb,
   rhsBatchDims: rb
 }?: DotDimensionNumbers): Array;
-type PaddingType = "VALID" | "SAME" | "SAME_LOWER" | [number, number][];
+type PaddingType = "VALID" | "SAME" | "SAME_LOWER" | Pair[];
 /**
  * General n-dimensional convolution operator, with optional dilation.
  *
@@ -1785,9 +2234,8 @@ declare function erfc(x: ArrayLike): Array;
  * forward or reverse-mode automatic differentiation.
  */
 declare function stopGradient(x: ArrayLike): Array;
-//# sourceMappingURL=lax.d.ts.map
 declare namespace nn_d_exports {
-  export { celu, elu, gelu, glu, identity, leakyRelu, logSigmoid, logSoftmax, logmeanexp, logsumexp, mish, oneHot, relu, relu6, sigmoid, silu, softSign, softmax, softplus, squareplus, standardize, swish };
+  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 };
 }
 /**
  * Rectified Linear Unit (ReLU) activation function:
@@ -1814,21 +2262,28 @@ declare function sigmoid(x: ArrayLike): Array;
  */
 declare function softplus(x: ArrayLike): Array;
 /**
- * Soft-sign activation function, computed element-wise:
- * `softsign(x) = x / (|x| + 1)`.
+ * @function
+ * Sparse plus function:
+ *
+ * - When `x <= -1`: `0`
+ * - When `-1 < x < 1`: `(x+1)**2 / 4`
+ * - When `x >= 1`: `x`
  */
-declare function softSign(x: ArrayLike): Array;
+declare const sparsePlus: OwnedFunction<(x: ArrayLike) => Array>;
 /**
  * @function
- * Sigmoid-weighted Linear Unit (SiLU) activation function, also known as
- * Swish, computed element-wise:
- * `silu(x) = x * sigmoid(x) = x / (1 + exp(-x))`.
+ * Sparse sigmoid activation function.
  *
- * `swish()` and `silu()` are both aliases for the same function.
- *
- * Reference: https://en.wikipedia.org/wiki/Swish_function
+ * - When `x <= -1`: `0`
+ * - When `-1 < x < 1`: `(x + 1) / 2`
+ * - When `x >= 1`: `1`
  */
-declare const silu: OwnedFunction<(x: ArrayLike) => Array>;
+declare const sparseSigmoid: OwnedFunction<(x: ArrayLike) => Array>;
+/**
+ * Soft-sign activation function, computed element-wise:
+ * `softsign(x) = x / (|x| + 1)`.
+ */
+declare function softSign(x: ArrayLike): Array;
 /**
  * @function
  * Sigmoid-weighted Linear Unit (SiLU) activation function, also known as
@@ -1839,7 +2294,7 @@ declare const silu: OwnedFunction<(x: ArrayLike) => Array>;
  *
  * Reference: https://en.wikipedia.org/wiki/Swish_function
  */
-declare const swish: OwnedFunction<(x: ArrayLike) => Array>;
+declare const silu: OwnedFunction<(x: ArrayLike) => Array>;
 /**
  * Log-sigmoid activation function, computed element-wise:
  * `log_sigmoid(x) = log(sigmoid(x)) = -log(1 + exp(-x))`.
@@ -1852,6 +2307,12 @@ declare function logSigmoid(x: ArrayLike): Array;
 declare const identity: (x: ArrayLike) => Array;
 /** Leaky rectified linear (ReLU) activation function */
 declare function leakyRelu(x: ArrayLike, negativeSlope?: ArrayLike): Array;
+/** Hard sigmoid activation function: `relu6(x+3)/6`. */
+declare function hardSigmoid(x: ArrayLike): Array;
+/** Hard SiLU (swish) activation function: `x * hardSigmoid(x)`. */
+declare function hardSilu(x: ArrayLike): Array;
+/** Hard tanh activation function: `clip(x, -1, 1)`. */
+declare function hardTanh(x: ArrayLike): Array;
 /**
  * Exponential linear unit activation function.
  *
@@ -1866,6 +2327,16 @@ declare function elu(x: ArrayLike, alpha?: ArrayLike): Array;
  * `celu(x) = x > 0 ? x : alpha * (exp(x/alpha) - 1)`
  */
 declare function celu(x: ArrayLike, alpha?: ArrayLike): Array;
+/**
+ * @function
+ * Scaled exponential linear unit activation.
+ *
+ * Computes the element-wise function:
+ * `selu(x) = lambda * (x > 0 ? x : alpha * (exp(x) - 1))`
+ *
+ * Where `alpha = 1.6732632423543772` and `lambda = 1.0507009873554805`.
+ */
+declare const selu: OwnedFunction<(x: ArrayLike) => Array>;
 /**
  * @function
  * Gaussion error linear unit (GELU) activation function.
@@ -1961,12 +2432,11 @@ declare function standardize(x: ArrayLike, axis?: Axis, opts?: {
  * ```
  */
 declare function oneHot(x: Array, numClasses: number): Array;
-//# sourceMappingURL=nn.d.ts.map
 declare namespace random_d_exports {
-  export { bernoulli, bits, exponential, key, 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. */
@@ -1986,11 +2456,50 @@ declare const uniform: OwnedFunction<(key: ArrayLike, shape?: number[] | undefin
  * and must be broadcastable to `shape`.
  */
 declare function bernoulli(key: Array, p?: ArrayLike, shape?: number[]): Array;
+/**
+ * @function
+ * Sample from a Cauchy distribution with location 0 and scale 1.
+ *
+ * Uses inverse transform sampling: `x = tan(π * (u - 0.5))` where u ~ Uniform(0, 1).
+ */
+declare const cauchy: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined) => Array>;
 /**
  * @function
  * Sample exponential random values according to `p(x) = exp(-x)`.
  */
 declare const exponential: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined) => Array>;
+/**
+ * @function
+ * Sample from a Gumbel distribution with location 0 and scale 1.
+ *
+ * Uses inverse transform sampling: `x = -log(-log(u))` where u ~ Uniform(0, 1).
+ */
+declare const gumbel: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined) => Array>;
+/**
+ * @function
+ * Sample from a Laplace distribution with location 0 and scale 1.
+ *
+ * Uses inverse transform sampling: the CDF is `F(x) = 0.5 + 0.5 * sign(x) * (1 - exp(-|x|))`.
+ * 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)`.
@@ -2000,7 +2509,6 @@ declare const exponential: OwnedFunction<(key: ArrayLike, shape?: number[] | und
  * bitwise identical to JAX.
  */
 declare const normal: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined) => Array>;
-//# sourceMappingURL=random.d.ts.map
 declare namespace scipy_special_d_exports {
   export { erf, erfc, logSoftmax, logit, logsumexp, softmax };
 }
@@ -2031,8 +2539,7 @@ declare const jacfwd: <F extends (x: Array) => Array>(f: F) => (...args: MapJsTr
  * Construct a Jaxpr by dynamically tracing a function with example inputs.
  */
 declare const makeJaxpr: <F extends (...args: any[]) => JsTree<Array>>(f: F) => (...args: Parameters<F>) => {
-  jaxpr: Jaxpr;
-  consts: Array[];
+  jaxpr: ClosedJaxpr;
   treedef: JsTreeDef;
 };
 /**
@@ -2080,11 +2587,6 @@ declare const valueAndGrad: <F extends (...args: any[]) => JsTree<Array>>(f: F)
  * Compute the Jacobian evaluated row-by-row by reverse-mode AD.
  */
 declare const jacrev: typeof jacfwd;
-/**
- * @function
- * Compute the Jacobian with reverse-mode AD. Alias for `jacrev()`.
- */
-declare const jacobian: <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()`.
  *
@@ -2106,8 +2608,5 @@ declare function blockUntilReady<T extends JsTree<any>>(x: T): Promise<T>;
  * default device.
  */
 declare function devicePut<T extends JsTree<any>>(x: T, device?: Device): Promise<MapJsTree<T, number | boolean, Array>>;
-//# sourceMappingURL=index.d.ts.map
-
 //#endregion
-export { Array, DType, type Device, Jaxpr, type JsTree, type JsTreeDef, type OwnedFunction, blockUntilReady, defaultDevice, devicePut, devices, grad, init, jacfwd, 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 };
-//# sourceMappingURL=index.d.ts.map
+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 };