@jax-js/jax 0.1.2 → 0.1.4

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
 /**
@@ -251,7 +250,7 @@ declare class AluExp implements FpHashable {
   /** Substitute variables in this AluExp to values. */
   substitute(variables: Record<string, AluExp>): AluExp;
   /** Reindex gid values in this expression as needed. */
-  reindexGids(gidMap: Map<number, number>): AluExp;
+  reindexGids(newGids: number[]): AluExp;
   get min(): number;
   get max(): number;
   /** Largest known integer that divides self. */
@@ -404,6 +403,52 @@ 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 questions. */
+  TriangularSolve = "TriangularSolve",
+  /** Cholesky decomposition of 2D positive semi-definite matrices. */
+  Cholesky = "Cholesky",
+}
+interface RoutineType {
+  inputShapes: number[][];
+  inputDtypes: DType[];
+  outputShapes: number[][];
+  outputDtypes: DType[];
+}
+//#endregion
 //#region src/backend.d.ts
 type Device = "cpu" | "wasm" | "webgpu";
 declare const devices: Device[];
@@ -441,9 +486,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 +503,69 @@ 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, diagonal, lstsq, matmul, matrixTranspose, outer, 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;
+/**
+ * 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;
 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, 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, 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 };
 }
 declare const float32 = DType.Float32;
 declare const int32 = DType.Int32;
@@ -587,6 +691,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. */
@@ -605,6 +723,13 @@ declare function argmin(a: ArrayLike, axis?: number, opts?: ReduceOpts): Array;
  * specified axis.
  */
 declare function argmax(a: ArrayLike, axis?: number, opts?: ReduceOpts): Array;
+/**
+ * Cumulative sum of elements along an axis.
+ *
+ * Currently this function is `O(n^2)`, we'll improve this later on with a
+ * two-phase parallel reduction algorithm.
+ */
+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;
 /**
@@ -650,12 +775,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.
  *
@@ -702,6 +844,22 @@ 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;
 /** 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;
@@ -773,6 +931,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.
  *
@@ -784,21 +946,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.
  *
@@ -815,8 +962,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). */
@@ -870,12 +1015,6 @@ 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. */
@@ -890,8 +1029,6 @@ 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;
 /** Round input to the nearest integer towards zero. */
 declare function trunc(x: ArrayLike): Array;
 /**
@@ -931,8 +1068,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>;
 /**
@@ -977,12 +1112,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.
  *
@@ -1009,6 +1138,10 @@ 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): 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). */
@@ -1022,7 +1155,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 };
 }
@@ -1073,13 +1205,18 @@ declare function dispose<Tree extends JsTree<Array>>(tree: Tree | null | undefin
 //#region src/frontend/convolution.d.ts
 /** Definition of a general dilated convolution. Should be valid on creation. */
 interface ConvParams {
+  vmapDims: number;
   strides: number[];
-  padding: [number, number][];
+  padding: Pair[];
   lhsDilation: number[];
   rhsDilation: number[];
 }
 /**
  * Check that the shapes and parameters passed to convolution are valid.
+ * Expected shapes of the lhs and rhs of the convolution are:
+ *
+ * - `lhsShape = [*vmapDims, batchSize, inChannels, spatialDims...]`
+ * - `rhsShape = [*vmapDims, outChannels, inChannels, kernelSize...]`
  *
  * If the check succeeds, returns the output shape.
  */
@@ -1151,9 +1288,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[];
@@ -1176,7 +1325,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",
@@ -1184,7 +1335,6 @@ declare enum Primitive {
   StopGradient = "stop_gradient",
   Cast = "cast",
   Bitcast = "bitcast",
-  RandomBits = "random_bits",
   Sin = "sin",
   Cos = "cos",
   Asin = "asin",
@@ -1194,8 +1344,6 @@ declare enum Primitive {
   Erf = "erf",
   Erfc = "erfc",
   Sqrt = "sqrt",
-  Min = "min",
-  Max = "max",
   Reduce = "reduce",
   Dot = "dot",
   // sum(x*y, axis=-1)
@@ -1205,14 +1353,23 @@ declare enum Primitive {
   PoolTranspose = "pool_transpose",
   Compare = "compare",
   Where = "where",
+  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
+  Jit = "jit",
 }
 interface PrimitiveParamsImpl extends Record<Primitive, Record<string, any>> {
   [Primitive.Cast]: {
@@ -1238,6 +1395,14 @@ interface PrimitiveParamsImpl extends Record<Primitive, Record<string, any>> {
   [Primitive.Compare]: {
     op: CompareOp;
   };
+  [Primitive.RandomBits]: {
+    shape: number[];
+    mode: "xor" | 0 | 1;
+  };
+  [Primitive.Gather]: {
+    axis: number[];
+    outDim: number;
+  };
   [Primitive.Transpose]: {
     perm: number[];
   };
@@ -1245,10 +1410,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[];
   };
@@ -1261,15 +1422,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>;
@@ -1437,7 +1597,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;
@@ -1456,6 +1616,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.
    *
@@ -1498,6 +1671,7 @@ declare class ShapedArray implements AbstractValue {
   constructor(shape: number[], dtype: DType, weakType: boolean);
   static fromAval(aval: AbstractValue): ShapedArray;
   get ndim(): number;
+  get size(): number;
   toString(): string;
   equals(other: ShapedArray): boolean;
 }
@@ -1515,12 +1689,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;
@@ -1552,7 +1726,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
@@ -1686,6 +1859,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.
  *
@@ -1699,8 +1887,71 @@ declare function linspace(start: number, stop: number, num?: number, endpoint?:
   dtype,
   device
 }?: DTypeAndDevice): Array;
+declare namespace lax_linalg_d_exports {
+  export { cholesky, 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;
+/**
+ * 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.
@@ -1731,7 +1982,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.
  *
@@ -1742,10 +1993,12 @@ type PaddingType = "VALID" | "SAME" | "SAME_LOWER" | [number, number][];
  */
 declare function convGeneralDilated(lhs: Array, rhs: Array, windowStrides: number[], padding: PaddingType, {
   lhsDilation,
-  rhsDilation
+  rhsDilation,
+  featureGroupCount
 }?: {
   lhsDilation?: number[];
   rhsDilation?: number[];
+  featureGroupCount?: number;
 }): Array;
 /** Convenience wrapper around `convGeneralDilated`. */
 declare function convWithGeneralPadding(lhs: Array, rhs: Array, windowStrides: number[], padding: PaddingType, lhsDilation?: number[], rhsDilation?: number[]): Array;
@@ -1769,9 +2022,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:
@@ -1798,21 +2050,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))`.
- *
- * `swish()` and `silu()` are both aliases for the same function.
+ * Sparse sigmoid activation 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
@@ -1823,7 +2082,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))`.
@@ -1836,6 +2095,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.
  *
@@ -1850,6 +2115,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.
@@ -1912,9 +2187,9 @@ declare function logSoftmax(x: ArrayLike, axis?: Axis): Array;
  *
  * Reference: https://en.wikipedia.org/wiki/LogSumExp
  */
-declare function logsumexp(x: ArrayLike, axis?: Axis): Array;
+declare function logsumexp(x: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
 /** Log-mean-exp reduction, like `jax.nn.logsumexp()` but subtracts `log(n)`. */
-declare function logmeanexp(x: ArrayLike, axis?: Axis): Array;
+declare function logmeanexp(x: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
 /**
  * Standardizes input to zero mean and unit variance.
  *
@@ -1945,9 +2220,8 @@ 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, normal, split, uniform };
 }
 /** Create a pseudo-random number generator (PRNG) key from 32-bit integer seed. */
 declare function key(seed: number): Array;
@@ -1970,11 +2244,33 @@ 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 random values according to `p(x) = 1/sqrt(2pi) * exp(-x^2/2)`.
@@ -1984,7 +2280,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 };
 }
@@ -2015,8 +2310,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;
 };
 /**
@@ -2064,11 +2358,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()`.
  *
@@ -2090,8 +2379,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 };