@jax-js/jax 0.1.5 → 0.1.7

package/dist/index.d.cts

@@ -541,1704 +541,1782 @@ declare class Executable<T = any> {
   source: Kernel | Routine, /** Extra data specific to the backend running this executable. */
   data: T);
 }
-declare namespace numpy_fft_d_exports {
-  export { ComplexPair, fft, ifft };
+declare namespace tree_d_exports {
+  export { JsTree, JsTreeDef, MapJsTree, NodeType, dispose, flatten, leaves, map, ref, structure, unflatten };
 }
-/**
- * A pair of arrays representing real and imaginary part `a + bj`. Both arrays
- * must have the same shape.
- */
-type ComplexPair = {
-  real: Array;
-  imag: Array;
+declare enum NodeType {
+  Array = "Array",
+  Object = "Object",
+  Leaf = "Leaf",
+}
+/** Analog to the JAX "pytree" object, but for JavaScript. */
+type JsTree<T> = T | JsTree<T>[] | {
+  [key: string]: JsTree<T>;
 };
-/**
- * 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 };
+type Same<X, Y> = (<T>() => T extends X ? 1 : 2) extends (<T>() => T extends Y ? 1 : 2) ? true : false;
+type MappedJsTree<T, A, B> = T extends A ? B : T extends Array ? T : T extends globalThis.Array<infer U> ? number extends T["length"] ? MapJsTree<U, A, B>[] : { [K in keyof T]: MapJsTree<T[K], A, B> } : { [K in keyof T]: MapJsTree<T[K], A, B> };
+/** @ignore Convert a subtype of JsTree<A> into a JsTree<B>, with the same structure. */
+type MapJsTree<T, A, B> = Same<A, B> extends true ? T : MappedJsTree<T, A, B>;
+/** Represents the structure of a JsTree. */
+declare class JsTreeDef {
+  readonly nodeType: NodeType;
+  readonly nodeMetadata: any;
+  readonly childTreedefs: JsTreeDef[];
+  static leaf: JsTreeDef;
+  constructor(nodeType: NodeType, nodeMetadata: any,
+  // Must be comparable with deepEqual.
+  childTreedefs: JsTreeDef[]);
+  /** Get the total number of leaves in the tree. */
+  get size(): number;
+  /** Returns a string representation of this tree definition. */
+  toString(root?: boolean): string;
+  /** Compare this tree definition with another. */
+  equals(other: JsTreeDef): boolean;
+}
+/** Flatten a structured object, returning the tree definition. */
+declare function flatten<T>(tree: JsTree<T>): [T[], JsTreeDef];
+/** Get the leaves of a tree. */
+declare function leaves<T>(tree: JsTree<T>): T[];
+/** Get the treedef for a tree. */
+declare function structure<T>(tree: JsTree<T>): JsTreeDef;
+/** Reconstruct a structured object from the flattened representation. */
+declare function unflatten<T>(treedef: JsTreeDef, leaves: Iterable<T>): JsTree<T>;
+/** Maps a multi-input function over pytree args to produce a new pytree. */
+declare function map<T, U, Tree extends JsTree<T>>(fn: (...args: T[]) => U, tree: Tree, ...rest: Tree[]): MapJsTree<Tree, T, U>;
+/** Take a reference of every array in a tree. */
+declare function ref<Tree extends JsTree<any>>(tree: Tree): Tree;
+/** Dispose every array in a tree. */
+declare function dispose<Tree extends JsTree<any>>(tree: Tree | null | undefined): void;
+//#endregion
+//#region src/frontend/convolution.d.ts
+/** Definition of a general dilated convolution. Should be valid on creation. */
+interface ConvParams {
+  vmapDims: number;
+  strides: number[];
+  padding: Pair[];
+  lhsDilation: number[];
+  rhsDilation: number[];
 }
 /**
- * Compute the Cholesky decomposition of a (batched) positive-definite matrix.
+ * Check that the shapes and parameters passed to convolution are valid.
+ * Expected shapes of the lhs and rhs of the convolution are:
  *
- * This is like `jax.lax.linalg.cholesky()`, except with an option to symmetrize
- * the input matrix, which is on by default.
+ * - `lhsShape = [*vmapDims, batchSize, inChannels, spatialDims...]`
+ * - `rhsShape = [*vmapDims, outChannels, inChannels, kernelSize...]`
+ *
+ * If the check succeeds, returns the output shape.
  */
-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;
+//#endregion
+//#region src/frontend/jaxpr.d.ts
 /**
- * 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.
+ * Function callback with an associated dispose() method.
  *
- * @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)`
+ * The dispose() method should be called to clean up any tracer resources needed
+ * by the function after the last time it is called.
  */
-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];
+type OwnedFunction<F extends Function> = F & {
+  dispose: () => void;
+};
+/** Variable in a Jaxpr expression. */
+declare class Var {
+  #private;
+  readonly id: number;
+  readonly aval: ShapedArray;
+  constructor(aval: ShapedArray);
+  toString(): string;
+}
+/** Literal in a Jaxpr expression. Currently, only scalars are supported. */
+declare class Lit {
+  readonly value: number;
+  readonly aval: ShapedArray;
+  get dtype(): DType;
+  constructor(aval: AbstractValue, value: number);
+}
+type Atom = Var | Lit;
+declare class VarPrinter {
+  #private;
+  names: Map<Var, string>;
+  name(v: Var): string;
+  nameType(v: Var): string;
+}
+/** A single statement / binding in a Jaxpr, in ANF form. */
+declare class JaxprEqn {
+  readonly primitive: Primitive;
+  readonly inputs: Atom[];
+  readonly params: Record<string, any>;
+  readonly outBinders: Var[];
+  constructor(primitive: Primitive, inputs: Atom[], params: Record<string, any>, outBinders: Var[]);
+  pprint(usedVars?: Set<Var>, vp?: VarPrinter): PPrint;
+  toString(): string;
+}
+/** Typed intermediate representation for traced computations. */
+declare class Jaxpr implements FpHashable {
+  #private;
+  readonly inBinders: Var[];
+  readonly eqns: JaxprEqn[];
+  readonly outs: Atom[];
+  constructor(inBinders: Var[], eqns: JaxprEqn[], outs: Atom[]);
+  pprint(): PPrint;
+  toString(): string;
+  /**
+   * Gets a hash of this Jaxpr.
+   *
+   * Var identity is not considered in the hash, so two Jaxprs with the same
+   * order of assignments and operators but different variable IDs will resolve
+   * to the same hash (and toString representation).
+   */
+  getHash(): bigint;
+  hash(state: FpHash): void;
+  /**
+   * Produce a simplified Jaxpr with basic optimizations applied.
+   *  - Trim away unused variables.
+   *  - Fold away *1, *0, or +0 operations against literals.
+   *  - Remove no-op movement operations.
+   */
+  simplify(): Jaxpr;
+  /** 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[];
+};
+//#endregion
+//#region src/frontend/core.d.ts
 /**
- * Solve a linear system of equations.
+ * Frontend primitive operations, which are lowered into Kernel objects before
+ * being dispatched to the backend.
  *
- * 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.
+ * Any operation between arrays can be described in these parts. This is also
+ * the set of primitives that can occur in Jaxpr programs, and the level at
+ * which transformations like vmap, grad, and jvp occur. They are loosely based
+ * on [XLA](https://openxla.org/xla/operation_semantics).
  *
- * @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)`.
+ * All n-ary operations support broadcasting, with NumPy semantics.
  */
-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, 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 enum Primitive {
+  Add = "add",
+  Mul = "mul",
+  Idiv = "idiv",
+  Mod = "mod",
+  // uses sign of numerator, C-style, matches JS but not Python
+  Min = "min",
+  Max = "max",
+  Neg = "neg",
+  Reciprocal = "reciprocal",
+  Floor = "floor",
+  Ceil = "ceil",
+  StopGradient = "stop_gradient",
+  Cast = "cast",
+  Bitcast = "bitcast",
+  Sin = "sin",
+  Cos = "cos",
+  Asin = "asin",
+  Atan = "atan",
+  Exp = "exp",
+  Log = "log",
+  Erf = "erf",
+  Erfc = "erfc",
+  Sqrt = "sqrt",
+  Reduce = "reduce",
+  Dot = "dot",
+  // sum(x*y, axis=-1)
+  Conv = "conv",
+  // see lax.conv_general_dilated
+  Pool = "pool",
+  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",
+  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",
 }
-declare const float32 = DType.Float32;
-declare const int32 = DType.Int32;
-declare const uint32 = DType.Uint32;
-declare const bool = DType.Bool;
-declare const float16 = DType.Float16;
-declare const float64 = DType.Float64;
-/** Euler's constant, `e = 2.7182818284590...` */
-declare const e: number;
-/** Euler-Mascheroni constant, `γ = 0.5772156649...` */
-declare const eulerGamma = 0.5772156649015329;
-/** Positive infinity. */
-declare const inf: number;
-/** Floating-point representation of NaN. */
-declare const nan: number;
-/** This is Pi, `π = 3.14159265358979...` */
-declare const pi: number;
-/** @function Element-wise addition, with broadcasting. */
-declare const add: (x: ArrayLike, y: ArrayLike) => Array;
-/** @function Element-wise multiplication, with broadcasting. */
-declare const multiply: (x: ArrayLike, y: ArrayLike) => Array;
-/** @function Numerical negative of every element of an array. */
-declare const negative: (x: ArrayLike) => Array;
-/** @function Calculate element-wise reciprocal of the input. This is `1/x`. */
-declare const reciprocal: (x: ArrayLike) => Array;
-/** @function Round input down to the nearest integer. */
-declare const floor: (x: ArrayLike) => Array;
-/** @function Round input up to the nearest integer. */
-declare const ceil: (x: ArrayLike) => Array;
-/** @function Element-wise sine function (takes radians). */
-declare const sin: (x: ArrayLike) => Array;
-/** @function Element-wise cosine function (takes radians). */
-declare const cos: (x: ArrayLike) => Array;
-/** @function Element-wise inverse sine function (inverse of sin). */
-declare const asin: (x: ArrayLike) => Array;
-/** @function Element-wise inverse tangent function (inverse of tan). */
-declare const atan: (x: ArrayLike) => Array;
-/** @function Calculate the exponential of all elements in the input array. */
-declare const exp: (x: ArrayLike) => Array;
-/** @function Calculate the natural logarithm of all elements in the input array. */
-declare const log: (x: ArrayLike) => Array;
-/** @function Calculate the square root of all elements in the input array. */
-declare const sqrt: (x: ArrayLike) => Array;
-/** @function Return element-wise minimum of the input arrays. */
-declare const minimum: (x: ArrayLike, y: ArrayLike) => Array;
-/** @function Return element-wise maximum of the input arrays. */
-declare const maximum: (x: ArrayLike, y: ArrayLike) => Array;
-/** @function Compare two arrays element-wise. */
-declare const greater: (x: ArrayLike, y: ArrayLike) => Array;
-/** @function Compare two arrays element-wise. */
-declare const less: (x: ArrayLike, y: ArrayLike) => Array;
-/** @function Compare two arrays element-wise. */
-declare const equal: (x: ArrayLike, y: ArrayLike) => Array;
-/** @function Compare two arrays element-wise. */
-declare const notEqual: (x: ArrayLike, y: ArrayLike) => Array;
-/** @function Compare two arrays element-wise. */
-declare const greaterEqual: (x: ArrayLike, y: ArrayLike) => Array;
-/** @function Compare two arrays element-wise. */
-declare const lessEqual: (x: ArrayLike, y: ArrayLike) => Array;
-/** @function Element-wise ternary operator, evaluates to `x` if cond else `y`. */
-declare const where: (cond: ArrayLike, x: ArrayLike, y: ArrayLike) => Array;
+interface PrimitiveParamsImpl extends Record<Primitive, Record<string, any>> {
+  [Primitive.Cast]: {
+    dtype: DType;
+  };
+  [Primitive.Bitcast]: {
+    dtype: DType;
+  };
+  [Primitive.Reduce]: {
+    op: AluOp;
+    axis: number[];
+  };
+  [Primitive.Conv]: ConvParams;
+  [Primitive.Pool]: {
+    window: number[];
+    strides: number[];
+  };
+  [Primitive.PoolTranspose]: {
+    inShape: number[];
+    window: number[];
+    strides: number[];
+  };
+  [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[];
+  };
+  [Primitive.Broadcast]: {
+    shape: number[];
+    axis: number[];
+  };
+  [Primitive.Reshape]: {
+    shape: number[];
+  };
+  [Primitive.Flip]: {
+    axis: number[];
+  };
+  [Primitive.Shrink]: {
+    slice: Pair[];
+  };
+  [Primitive.Pad]: {
+    width: Pair[];
+  };
+  [Primitive.TriangularSolve]: {
+    unitDiagonal: boolean;
+  };
+  [Primitive.Jit]: {
+    name: string;
+    jaxpr: Jaxpr;
+    numConsts: number;
+  };
+}
+/** Type of parameters taken by each primitive. */
+type PrimitiveParams<T extends Primitive> = T extends keyof PrimitiveParamsImpl ? PrimitiveParamsImpl[T] : Record<string, never>;
+declare enum CompareOp {
+  Less = "less",
+  Equal = "equal",
+  NotEqual = "not_equal",
+  LessEqual = "less_equal",
+}
+/** @inline */
+type Axis = number | number[] | null;
+/** @inline */
+type ReduceOpts = {
+  keepdims?: boolean;
+};
+type MainTrace = {
+  level: number;
+  traceType: new (main: MainTrace) => Trace;
+  globalData: any | null;
+};
 /**
- * @function
- * Permute the dimensions of an array. Defaults to reversing the axis order.
+ * Push an interpreter onto the trace stack. Use this like:
+ * `using main = newMain(...);`
  */
-declare const transpose: (x: ArrayLike, perm?: number[]) => Array;
+
+type TracerValue = Tracer | number | boolean;
+declare abstract class Trace {
+  readonly main: MainTrace;
+  constructor(main: MainTrace);
+  abstract pure(val: TracerValue): Tracer;
+  abstract lift(val: Tracer): Tracer;
+  abstract processPrimitive<P extends Primitive>(primitive: P, tracers: Tracer[], params: PrimitiveParams<P>): Tracer[];
+}
+/** Internal representation of an array value. */
+interface AbstractValue {
+  /** Shape of the array. Must be a static tuple of non-negative dimensions. */
+  shape: number[];
+  /** Concrete data type of array elements. */
+  dtype: DType;
+  /**
+   * Arrays created from JavaScript numbers (e.g., `np.array(3)`) are created as
+   * _weakly typed_ unless a dtype is explicitly specified.
+   *
+   * Weakly typed values will automatically cast to the data type of other
+   * arrays when used as an operand as an expression. This property only affects
+   * how they promote in type casting; their memory layout is still determined
+   * by the actual `dtype` field.
+   *
+   * ```ts
+   * const x = np.array(3); // weakType = true, dtype = float32
+   * const y = np.array([1, 2], { dtype: np.int32 }); // weakType = false, dtype = int32
+   * const z = x.add(y); // z has dtype int32 because x is weakly typed
+   * ```
+   *
+   * Weak types are present in JIT programs in their spec (e.g., Jaxpr inputs
+   * and outputs can be weakly typed) form. But they're solely a frontend
+   * concept. Backends are not aware of weak types.
+   */
+  weakType: boolean;
+}
 /**
- * @function
- * Give a new shape to an array without changing its data.
+ * Broadcast shapes and promote types with casting for two avals.
  *
- * One shape dimension can be -1. In this case, the value is inferred from the
- * length of the array and remaining dimensions.
+ * This implements the weak type behavior described in `promoteTypes()`, but not
+ * implemented in that function as `weakType` is not passed.
  */
-declare const reshape: (x: ArrayLike, shape: number[]) => Array;
-/**
- * @function
- * Move axes of an array to new positions. Other axes retain original order.
- */
-declare const moveaxis: (x: ArrayLike, src: number, dst: number) => Array;
+
+declare abstract class Tracer {
+  /** @ignore */
+  readonly _trace: Trace;
+  constructor(trace: Trace);
+  abstract get aval(): AbstractValue;
+  abstract toString(): string;
+  /**
+   * Access an array by reference, incrementing the reference count.
+   *
+   * jax-js handles freeing arrays by using "move" semantics, like in Rust/C++.
+   * Whenever you pass an array into a function, that function should consume
+   * the array, and it will no longer be usable. For example, if you had:
+   *
+   * ```
+   * const x = np.array([1, 2, 3]);
+   * const y = np.add(x, x);
+   * ```
+   *
+   * The second line does not work because the first parameter consumes `x`, and
+   * then the second parameter will already have been freed / disposed.
+   *
+   * To fix this, you can write:
+   *
+   * ```
+   * const y = np.add(x.ref, x);
+   * ```
+   *
+   * Under the hood, every access to `.ref` increments the internal reference
+   * count of the array. The reference count starts at 1. When it hits 0, the
+   * memory behind the array is freed.
+   */
+  abstract get ref(): this;
+  /**
+   * Manually decrement the reference count of the array.
+   *
+   * Arrays are created with reference count 1. Whenever it is used as argument
+   * to a function or other operation, it is disposed (i.e., reference count
+   * decreases by 1) automatically. Whenever a `.ref` is created, the reference
+   * count increases.
+   *
+   * You generally don't need to call this function directly since arrays are
+   * automatically disposed after being passed into an operation. One common
+   * exception is when writing a function and ignoring one of its arguments. In
+   * that case, by convention you should dispose of that argument manually.
+   *
+   * ```
+   * function myCustomOperation(a: np.Array, b: np.Array) {
+   *   b.dispose(); // Needed to satisfy "move" rules.
+   *   return a.add(1);
+   * }
+   * ```
+   */
+  abstract dispose(): void;
+  /** The shape of the array. */
+  get shape(): number[];
+  /** The total number of elements in the array. */
+  get size(): number;
+  /** The dtype of elements stored in the array. */
+  get dtype(): DType;
+  /**
+   * Whether the array is weakly typed.
+   *
+   * Weakly typed arrays will cast to the dtype of the other operand. See
+   * `promoteTypes()` for details.
+   */
+  get weakType(): boolean;
+  /** The number of dimensions of the array. */
+  get ndim(): number;
+  /** @ignore */
+  fullLower(): 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;
+  notEqual(other: this | TracerValue): this;
+  greaterEqual(other: this | TracerValue): this;
+  lessEqual(other: this | TracerValue): this;
+  /** Sum of the elements of the array over a given axis, or axes. */
+  sum(axis?: Axis, opts?: ReduceOpts): this;
+  /** Product of the array elements over a given axis. */
+  prod(axis?: Axis, opts?: ReduceOpts): this;
+  /** Compute the average of the array elements along the specified axis. */
+  mean(axis?: Axis, opts?: ReduceOpts): this;
+  /** Minimum of the elements of the array along a given axis. */
+  min(axis?: Axis, opts?: ReduceOpts): this;
+  /** Maximum of the elements of the array along a given axis. */
+  max(axis?: Axis, opts?: ReduceOpts): this;
+  /** Test whether all array elements along a given axis evaluate to true. */
+  all(axis?: Axis, opts?: ReduceOpts): this;
+  /** Test whether any array element along a given axis evaluates to true. */
+  any(axis?: Axis, opts?: ReduceOpts): this;
+  /** Permute the dimensions of an array. Defaults to reversing the axis order. */
+  transpose(perm?: number[]): this;
+  /**
+   * Give a new shape to an array without changing its data.
+   *
+   * One shape dimension can be -1. In this case, the value is inferred from the
+   * length of the array and remaining dimensions.
+   */
+  reshape(shape: number | number[]): this;
+  /** Copy the array and cast to a specified dtype. */
+  astype(dtype: DType): this;
+  /** Subtract an array from this one. */
+  sub(other: this | TracerValue): this;
+  /** Divide an array by this one. */
+  div(other: this | TracerValue): this;
+  /** 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;
+  /** Flatten the array without changing its data. */
+  ravel(): this;
+  /**
+   * Iterate over the first dimension of this array, returning slices.
+   *
+   * This can be used to destructure arrays. For example:
+   *
+   * ```js
+   * let x = np.array([[1, 2], [3, 4]]);
+   * let [a, b] = x;
+   * console.log(a.js()); // [1, 2]
+   * console.log(b.js()); // [3, 4]
+   * ```
+   */
+  [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.
+   *
+   * This is the equivalent of slicing in Python, e.g. `x[1:3, 2, :, None]`. To
+   * mimic this in JavaScript, we would write:
+   *
+   * ```js
+   * x.slice([1, 3], 2, [], null);
+   * ```
+   *
+   * The `slice` method accepts a variable number of arguments, each of which
+   * can be a number, an empty array, a single-element array, a two-element
+   * array, or `null`. The arguments are interpreted as follows:
+   *
+   * - A number `n` means to access the `n`-th element along that axis, removing
+   *   that axis from the resulting shape.
+   * - An empty array `[]` means to keep that axis as-is, like `:` in Python.
+   * - A single-element array `[i]` means to start slicing from index `i`
+   *   (inclusive) to the end of the axis, like `x[i:]`.
+   * - A two-element array `[i, j]` means to slice from index `i` (inclusive)
+   *   to index `j` (exclusive), like `x[i:j]`.
+   * - `null` means to add a new axis at that position, like `np.newaxis`.
+   *
+   * Like in Python, negative indices are supported, which count from the end of
+   * the axis. For example, `-1` means the last element.
+   *
+   * Strided slices are not yet implemented, so you cannot write `x[::2]` or
+   * similar.
+   *
+   * Advanced indexing by integer arrays is also supported. This translates to
+   * the "gather" primitive, and it allows you to access specific elements of
+   * the array by integer indices stored in another array.
+   */
+  slice(...index: (number | [] | [number] | Pair | null | Tracer)[]): this;
+}
+declare class ShapedArray implements AbstractValue {
+  readonly shape: number[];
+  readonly dtype: DType;
+  readonly weakType: boolean;
+  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;
+}
+//#endregion
+//#region src/frontend/array.d.ts
+type ArrayLike = Array | number | boolean;
+/** Version of pureArray with fudged types. */
+
 /**
- * @function
- * Add padding (zeros) to an array.
+ * An executable operation that will be dispatched to the backend.
  *
- * The `width` argument is either an integer or pair of integers, in which case
- * all axes are padded with the same width. Or if it is an array of pairs, each
- * pair specifies the padding for its corresponding axis.
- */
-declare const pad: (x: ArrayLike, width: number | Pair | Pair[]) => Array;
-/**
- * @function
- * Return the number of dimensions of an array. Does not consume array reference.
- */
-declare const ndim: (x: ArrayLike) => number;
-/** @function Return the shape of an array. Does not consume array reference. */
-declare const shape$1: (x: ArrayLike) => number[];
-/**
- * @function
- * Return an array of zeros with the same shape and type as a given array.
- */
-declare const zerosLike: (a: ArrayLike, dtype?: DType) => Array;
-/**
- * @function
- * Return an array of ones with the same shape and type as a given array.
- */
-declare const onesLike: (a: ArrayLike, dtype?: DType) => Array;
-/**
- * @function
- * Return a full array with the same shape and type as a given array.
- */
-declare const fullLike: (a: ArrayLike, fillValue: number | boolean | Array, dtype?: DType) => Array;
-/**
- * Return the number of elements in an array, optionally along an axis.
- * Does not consume array reference.
+ * This holds a reference to all input buffers used in the operation. After the
+ * operation is dispatched, the references should be released.
  */
-declare function size(a: ArrayLike, axis?: number): number;
-/** Convert an array to a specified dtype. */
-declare function astype(a: ArrayLike, dtype: DType): Array;
-/** Sum of the elements of the array over a given axis, or axes. */
-declare function sum(a: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
-/** Product of the array elements over a given axis. */
-declare function prod(a: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
-/** Return the minimum of array elements along a given axis. */
-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;
+declare class PendingExecute {
+  #private;
+  readonly backend: Backend;
+  readonly source: Kernel | Routine;
+  readonly inputs: Slot[];
+  readonly outputs: Slot[];
+  prepared: Executable | null;
+  submitted: boolean;
+  constructor(backend: Backend, source: Kernel | Routine, inputs: Slot[], outputs: Slot[]);
+  updateRc(delta: number): void;
+  prepare(): Promise<void>;
+  prepareSync(): void;
+  submit(): void;
+}
+/** @inline */
+type DTypeAndDevice = {
+  dtype?: DType;
+  device?: Device;
+};
+type ArrayConstructorArgs = {
+  source: AluExp | Slot;
+  st: ShapeTracker;
+  dtype: DType;
+  weakType: boolean;
+  backend: Backend;
+  committed: boolean;
+  pending?: Iterable<PendingExecute>;
+};
 /**
- * Test whether all array elements along a given axis evaluate to True.
+ * A multidimensional numeric array with data stored on CPU or GPU.
  *
- * Returns a boolean array with the same shape as `a` with the specified axis
- * removed. If axis is None, returns a scalar.
+ * This is the library's core data type. Equivalent to `jax.Array` from JAX, or
+ * `torch.Tensor`.
+ *
+ * Not to be confused with the JavaScript "Array" constructor. Avoid importing
+ * this into your code's namespace if you're already using the JavaScript
+ * "Array" type by name.
  */
-declare function all(a: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
+declare class Array extends Tracer {
+  #private;
+  /**
+   * @ignore
+   * Constructs an array from source, shape and backend. Note that if the source
+   * is a backend `Slot`, this constructor _takes ownership_ of the slot. It
+   * will be freed when the array is disposed.
+   */
+  constructor(args: ArrayConstructorArgs);
+  /** @ignore */
+  get aval(): ShapedArray;
+  /** Return a simple string representation of the array's dimensions. */
+  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.
+   *
+   * This only works for scalars (0-dimensional arrays). It lets you get values
+   * "out" of the JAX system. For instance, if `x = np.array(5)`, then you can
+   * evaluate `x + 1` and `x ** 2` to get `6` and `25`, respectively.
+   *
+   * This method is also called for `==` equality.
+   */
+  [Symbol.toPrimitive](): any;
+  /** Realize the array and return it as data. */
+  data(): Promise<DataArray>;
+  /**
+   * Wait for this array to finish evaluation.
+   *
+   * Operations and data loading in jax-js are lazy, so this function ensures
+   * that pending operations are dispatched and fully executed before it
+   * returns.
+   *
+   * If you are mapping from `data()` or `dataSync()`, it will also trigger
+   * dispatch of operations as well.
+   *
+   * **Note:** `jax.blockUntilReady()` is a higher-level API, it calls this
+   * asynchronously for multiple arrays.
+   */
+  blockUntilReady(): Promise<Array>;
+  /**
+   * Realize the array and return it as data. This is a sync variant and not
+   * recommended for performance reasons, as it will block rendering.
+   */
+  dataSync(): DataArray;
+  /**
+   * Convert this array into a JavaScript object.
+   *
+   * This is a blocking operation that will compile all of the shaders and wait
+   * for execution to complete, synchronously. No other JavaScript code on the
+   * site will be run during shader execution.
+   *
+   * To avoid blocking, prefer `jsAsync()` when possible.
+   */
+  js(): any;
+  /** Convert this array into a JavaScript object, asynchronously. */
+  jsAsync(): Promise<any>;
+  /**
+   * Copy an element of an array to a numeric scalar and return it.
+   *
+   * Throws an error if the array does not have a single element. The array must
+   * either be rank-0, or all dimensions of the shape are 1.
+   */
+  item(): number;
+  /** @private Internal plumbing method for Array / Tracer ops. */
+  static _implRules(): typeof implRules;
+  /** @private */
+  _realizeSource(): number;
+  /** @private Put this array on a new backend, asynchronously. */
+  _put(backend: Backend): Promise<Array>;
+  /** @private Put this array on a new backend, synchronously. */
+  _putSync(backend: Backend): Array;
+}
+/** Constructor for creating a new array from data. */
+declare function array(values: Array | DataArray | RecursiveArray<number> | RecursiveArray<boolean>, {
+  shape,
+  dtype,
+  device
+}?: {
+  shape?: number[];
+} & DTypeAndDevice): Array;
+/** If x is a value, lift it into an array, otherwise leave it be. */
+
+type ImplRule<P extends Primitive> = (tracers: Array[], params: PrimitiveParams<P>) => Array[];
+declare const implRules: { [P in Primitive]: ImplRule<P> };
+/** Return a new array of given shape and type, filled with zeros. */
+declare function zeros(shape: number[], {
+  dtype,
+  device
+}?: DTypeAndDevice): Array;
+/** Return a new array of given shape and type, filled with ones. */
+declare function ones(shape: number[], {
+  dtype,
+  device
+}?: DTypeAndDevice): Array;
+/** Return a new array of given shape and type, filled with `fill_value`. */
+declare function full(shape: number[], fillValue: number | boolean | Array, {
+  dtype,
+  device
+}?: DTypeAndDevice): Array;
 /**
- * Test whether any array element along a given axis evaluates to True.
+ * Create an identity matrix.
  *
- * Returns a boolean array with the same shape as `a` with the specified axis
- * removed. If axis is None, returns a scalar.
+ * If numCols is not provided, it defaults to numRows, i.e., a square identity
+ * matrix with ones on the diagonal.
  */
-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. */
-declare function mean(a: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
+declare function eye(numRows: number, numCols?: number, {
+  dtype,
+  device
+}?: DTypeAndDevice): Array;
+/** Return the identity matrix, with ones on the main diagonal. */
+declare function identity$1(n: number, {
+  dtype,
+  device
+}?: DTypeAndDevice): Array;
 /**
- * Returns the indices of the minimum values along an axis.
+ * Return evenly spaced values within a given interval.
  *
- * By default, index is into the flatted array, otherwise it is along the
- * specified axis.
+ * This can be called with a varying number of arguments, just like the range()
+ * builtin function in Python.
+ *
+ * - `arange(stop)` is equivalent to `arange(0, stop, 1)`.
+ * - `arange(start, stop)` is equivalent to `arange(start, stop, 1)`.
+ * - `arange(start, stop, step)` creates an array starting at `start`, ending
+ *   before `stop`, with a step size of `step`.
+ *
+ * Defaults to an integer data type. This can produce unintended results when
+ * using a non-integer step, so prefer linspace() in those cases.
  */
-declare function argmin(a: ArrayLike, axis?: number, opts?: ReduceOpts): Array;
+declare function arange(start: number, stop?: number, step?: number, {
+  dtype,
+  device
+}?: DTypeAndDevice): Array;
 /**
- * Returns the indices of the maximum values along an axis.
+ * Return an array with ones on and below the diagonal and zeros elsewhere.
  *
- * By default, index is into the flatted array, otherwise it is along the
- * specified axis.
+ * 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 argmax(a: ArrayLike, axis?: number, opts?: ReduceOpts): Array;
+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;
 /**
- * Cumulative sum of elements along an axis.
+ * Return evenly spaced numbers over a specified interval.
  *
- * Currently this function is `O(n^2)`, we'll improve this later on with a
- * two-phase parallel reduction algorithm.
+ * Returns _num_ evenly spaced samples, calculated over the interval
+ * [`start`, `stop`]. The endpoint `stop` is included in the result by default,
+ * but this is controlled by the `endpoint` parameter.
+ *
+ * The default data type is Float32. Use arange() for integer steps.
  */
-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;
+declare function linspace(start: number, stop: number, num?: number, endpoint?: boolean, {
+  dtype,
+  device
+}?: DTypeAndDevice): Array;
 /**
- * Split an array into multiple sub-arrays along an axis.
+ * Return numbers spaced evenly on a log scale.
  *
- * @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.
+ * 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 split$1(a: ArrayLike, indicesOrSections: number | number[], axis?: number): Array[];
+declare function logspace(start: number, stop: number, num?: number, endpoint?: boolean, base?: number, {
+  dtype,
+  device
+}?: DTypeAndDevice): Array;
+//#endregion
+//#region src/frontend/linearize.d.ts
+/** @inline */
+type GradOpts = {
+  /**
+   * Integer or sequence of integers. Specifies which positional argument(s) to
+   * differentiate with respect to.
+   *
+   * Defaults to `0` (the first argument).
+   */
+  argnums?: number | number[];
+  /**
+   * The input function returns a pair of `[out, aux]` including an auxiliary
+   * value. This `aux` is not differentiated, but is returned alongside the
+   * gradient when evaluating the function.
+   */
+  hasAux?: boolean;
+};
+declare namespace lax_linalg_d_exports {
+  export { cholesky$1 as cholesky, lu, triangularSolve };
+}
 /**
- * Join a sequence of arrays along an existing axis.
+ * Compute the Cholesky decomposition of a symmetric positive-definite matrix.
  *
- * The arrays must have the same shape, except in the dimension corresponding to
- * `axis` (the first, by default).
+ * The Cholesky decomposition of a matrix `A` is:
  *
- * No scalars can be passed to this function, as the axis is then ambiguous.
+ * - 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 concatenate(xs: Array[], axis?: number): Array;
+declare function cholesky$1(a: ArrayLike, {
+  upper
+}?: {
+  upper?: boolean;
+}): Array;
 /**
- * Join a sequence of arrays along a new axis.
+ * LU decomposition with partial pivoting.
  *
- * The `axis` parameter specifies the index of the new axis in the dimensions of
- * the result. For example, if `axis=0` it will be the first dimension and if
- * `axis=-1` it will be the last dimension.
+ * 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.
  *
- * All shapes must have the same shape.
+ * @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 stack(xs: ArrayLike[], axis?: number): Array;
+declare function lu(x: ArrayLike): [Array, Array, Array];
 /**
- * Horizontally stack arrays. Inputs are promoted to rank at least 1, then
- * concatenated along axis 1 (if rank-2 or higher) or 0 (if rank-1).
+ * 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 hstack(xs: ArrayLike[]): Array;
+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, convTranspose, convWithGeneralPadding, dot$1 as dot, erf, erfc, lax_linalg_d_exports as linalg, reduceWindow, stopGradient };
+}
 /**
- * Vertically stack arrays. Inputs are promoted to rank at least 2, then
- * concatenated along axis 0.
+ * Dimension numbers for general `dot()` primitive.
+ *
+ * Contracting dimensions act as a tensor contraction (reduction) along the
+ * given axis. They must be the same size in both operands. Batch dimensions
+ * are treated as vectorized, leading batch dimensions.
+ *
+ * The return value has a shape where the first dimensions are shared batch
+ * dimensions, followed by `lhs` non-contracting dimensions, followed by
+ * `rhs` non-contracting dimensions.
  */
-declare function vstack(xs: ArrayLike[]): Array;
+type DotDimensionNumbers = {
+  lhsContractingDims?: number[];
+  rhsContractingDims?: number[];
+  lhsBatchDims?: number[];
+  rhsBatchDims?: number[];
+};
 /**
- * Stack arrays depth-wise. Inputs are promoted to rank at least 3, then
- * concatenated along axis 2.
+ * General dot product/contraction operator.
+ *
+ * Prefer higher-level functions like `jax.numpy.dot()`, `jax.numpy.matmul()`,
+ * `jax.numpy.tensordot(), and `jax.numpy.einsum()` where possible.
  */
-declare function dstack(xs: ArrayLike[]): Array;
+declare function dot$1(lhs: Array, rhs: Array, {
+  lhsContractingDims: lc,
+  rhsContractingDims: rc,
+  lhsBatchDims: lb,
+  rhsBatchDims: rb
+}?: DotDimensionNumbers): Array;
+type PaddingType = "VALID" | "SAME" | "SAME_LOWER" | Pair[];
 /**
- * Stack arrays column-wise. Inputs are promoted to rank at least 2, then
- * concatenated along axis 1.
+ * General n-dimensional convolution operator, with optional dilation.
+ *
+ * The semantics of this operation mimic the `jax.lax.conv_general_dilated`
+ * function in JAX, which wraps XLA's general convolution operator.
+ *
+ * @param lhs - Input tensor; shape `[N, C_in, ...xs]`
+ * @param rhs - Convolution kernel; shape `[C_out, C_in / G, ...ks]`
+ * @param windowStrides - Strides for each spatial dimension
+ * @param padding - Padding for each spatial dimension, or a string
+ *   (`"VALID"`, `"SAME"`, or `"SAME_LOWER"`)
  */
-declare function columnStack(xs: ArrayLike[]): Array;
-/** Flip an array vertically (axis=0). */
-declare function flipud(x: ArrayLike): Array;
-/** Flip an array horizontally (axis=1). */
-declare function fliplr(x: ArrayLike): 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;
+declare function convGeneralDilated(lhs: Array, rhs: Array, windowStrides: number[], padding: PaddingType, {
+  lhsDilation,
+  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;
+/** Convenience wrapper around `convGeneralDilated`. */
+declare function conv(lhs: Array, rhs: Array, windowStrides: number[], padding: PaddingType): Array;
 /**
- * Expand the shape of an array by inserting new axes of length 1.
+ * Convenience wrapper for calculating the N-d convolution "transpose".
  *
- * @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.
+ * This function directly calculates a fractionally strided conv rather than
+ * indirectly calculating the gradient (transpose) of a forward convolution.
+ * It is equivalent to the JAX version, except:
  *
- * @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]
- * ```
+ * - The `use_consistent_padding` option is not available. We only have the
+ *   consistent padding case (JAX version >0.8.4).
+ * - The order of dimensions matches `lax.conv_general_dilated`.
+ *
+ * Unlike PyTorch/TensorFlow, by default we don't reverse the kernel's spatial
+ * dimensions or the `(C_out, C_in)` axis order. To get this behavior, set
+ * `transposeKernel` to true.
+ *
+ * @param lhs - Input tensor; shape `[N, C_in, ...xs]`
+ * @param rhs - Convolution kernel; shape `[C_out, C_in, ...ks]`
+ * @param strides - Sequence of n integers, sets fractional stride
+ * @param padding - Apply padding of `dilation * (kernel_size - 1) - padding` to
+ *   each side of the input, so it acts like gradient of `conv()`
+ * @param rhsDilation - Atrous dilation for the kernel
+ * @param transposeKernel - Flip spatial axes and swap the input/output channels
+ *   of the kernel; its shape should be `[C_in, C_out, ...ks]`
  */
-declare function expandDims(a: ArrayLike, axis: number | number[]): Array;
+declare function convTranspose(lhs: Array, rhs: Array, strides: number[], padding: PaddingType, {
+  rhsDilation,
+  transposeKernel
+}?: {
+  rhsDilation?: number[];
+  transposeKernel?: boolean;
+}): Array;
+/** Reduce a computation over padded windows. */
+declare function reduceWindow(operand: Array, computation: (x: Array) => Array, windowDimensions: number[], windowStrides?: number[]): Array;
+/** The error function: `erf(x) = 2/sqrt(pi) * int[0..x] exp(-t^2) dt`. */
+declare function erf(x: ArrayLike): Array;
 /**
- * Repeat each element of an array after themselves.
+ * The complementary error function: `erfc(x) = 1 - erf(x)`.
  *
- * If no axis is provided, use the flattened input array, and return a flat
- * output array.
+ * This function is more accurate than `1 - erf(x)` for large values of `x`,
+ * where `erf(x)` is very close to 1.
  */
-declare function repeat(a: ArrayLike, repeats: number, axis?: number): Array;
+declare function erfc(x: ArrayLike): Array;
 /**
- * Construct an array by repeating A the number of times given by reps.
+ * Stops gradient computation.
  *
- * If `A` is an array of shape `(d1, d2, ..., dn)` and `reps` is a sequence of
- * integers, the resulting array will have a shape of `(reps[0] * d1,
- * reps[1] * d2, ..., reps[n] * dn)`, with `A` tiled along each dimension.
+ * Behaves as the identity function but prevents the flow of gradients during
+ * forward or reverse-mode automatic differentiation.
  */
-declare function tile(a: ArrayLike, reps: number | number[]): Array;
+declare function stopGradient(x: ArrayLike): Array;
+declare namespace numpy_fft_d_exports {
+  export { ComplexPair, fft, ifft };
+}
 /**
- * Broadcast an array to a shape, with NumPy-style broadcasing rules.
- *
- * In other words, this lets you append axes to the left, and/or expand
- * dimensions where the shape is 1.
+ * A pair of arrays representing real and imaginary part `a + bj`. Both arrays
+ * must have the same shape.
  */
-declare function broadcastTo(a: ArrayLike, shape: number[]): Array;
-/** Broadcast input shapes to a common output shape. */
-declare function broadcastShapes(...shapes: number[][]): number[];
-/** Broadcast arrays to a common shape. */
-declare function broadcastArrays(...arrays: ArrayLike[]): Array[];
+type ComplexPair = {
+  real: Array;
+  imag: Array;
+};
 /**
- * Return specified diagonals.
- *
- * If a is 2D, return the diagonal of the array with the given offset. If a is
- * 3D or higher, compute diagonals along the two given axes (default: 0, 1).
+ * Compute a one-dimensional discrete Fourier transform.
  *
- * This returns a view over the existing array. The shape of the resulting array
- * is determined by removing the two axes along which the diagonal is taken,
- * then appending a new axis to the right with holding the diagonals.
+ * Currently, the size of the axis must be a power of two.
  */
-declare function diagonal(a: ArrayLike, offset?: number, axis1?: number, axis2?: number): Array;
+declare function fft(a: ComplexPair, axis?: number): ComplexPair;
 /**
- * Extract a diagonal or construct a diagonal array.
+ * Compute a one-dimensional inverse discrete Fourier transform.
  *
- * If v is a 2D array, return the k-th diagonal of v (as a view). If v is a 1D
- * array, return a 2D array with v on the k-th diagonal.
+ * Currently, the size of the axis must be a power of two.
  */
-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;
+declare function ifft(a: ComplexPair, axis?: number): ComplexPair;
+declare namespace numpy_linalg_d_exports {
+  export { cholesky, det, diagonal, inv, lstsq, matmul, matrixPower, matrixTranspose, outer, slogdet, solve, tensordot, trace, vecdot };
+}
 /**
- * Return a sorted copy of an array.
+ * Compute the Cholesky decomposition of a (batched) positive-definite matrix.
  *
- * 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.
+ * This is like `jax.lax.linalg.cholesky()`, except with an option to symmetrize
+ * the input matrix, which is on by default.
  */
-declare function sort(a: ArrayLike, axis?: number): Array;
+declare function cholesky(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 indices that would sort an array. This may be an unstable sorting
- * algorithm; it need not preserve order of indices in ties.
+ * Return the least-squares solution to a linear equation.
  *
- * Returns an array of `int32` indices.
+ * For overdetermined systems, this finds the `x` that minimizes `norm(ax - b)`.
+ * For underdetermined systems, this finds the minimum-norm solution for `x`.
  *
- * The array is sorted along a specified axis (the last by default).
+ * 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 argsort(a: ArrayLike, axis?: number): Array;
+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];
 /**
- * Take elements from an array along an axis.
+ * Solve a linear system of equations.
  *
- * This is equivalent to advanced indexing with integer indices over that
- * numbered axis. By default, the flattened array is used.
+ * 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 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;
-  atol?: number;
-}): boolean;
-/** Matrix product of two arrays. */
-declare function matmul(x: ArrayLike, y: ArrayLike): Array;
-/** Dot product of two arrays. */
-declare function dot$1(x: ArrayLike, y: ArrayLike): Array;
+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, divmod, dot, dstack, e, einsum, equal, eulerGamma, exp, exp2, expandDims, expm1, eye, numpy_fft_d_exports as fft, finfo, flip, fliplr, flipud, float16, float32, float64, floor, floorDivide, fmod, frexp, full, fullLike, greater, greaterEqual, hamming, hann, heaviside, hstack, hypot, identity$1 as identity, iinfo, inf, inner, int32, isfinite, isinf, isnan, isneginf, isposinf, ldexp, less, lessEqual, numpy_linalg_d_exports as linalg, linspace, log, log10, log1p, log2, logspace, matmul, matrixTranspose, max, maximum, mean, meshgrid, min, minimum, moveaxis, multiply, nan, nanToNum, ndim, negative, notEqual, ones, onesLike, outer, pad, transpose as permuteDims, pi, positive, power as pow, power, prod, promoteTypes, ptp, rad2deg, radians, ravel, reciprocal, remainder, repeat, reshape, shape$1 as shape, sign, sin, sinc, sinh, size, sort, split$1 as split, sqrt, square, squeeze, stack, std, subtract, sum, swapaxes, take, tan, tanh, tensordot, tile, trace, transpose, tri, tril, triu, trueDivide, trunc, uint32, var_, vdot, vecdot, vstack, where, zeros, zerosLike };
+}
+declare const float32 = DType.Float32;
+declare const int32 = DType.Int32;
+declare const uint32 = DType.Uint32;
+declare const bool = DType.Bool;
+declare const float16 = DType.Float16;
+declare const float64 = DType.Float64;
+/** Euler's constant, `e = 2.7182818284590...` */
+declare const e: number;
+/** Euler-Mascheroni constant, `γ = 0.5772156649...` */
+declare const eulerGamma = 0.5772156649015329;
+/** Positive infinity. */
+declare const inf: number;
+/** Floating-point representation of NaN. */
+declare const nan: number;
+/** This is Pi, `π = 3.14159265358979...` */
+declare const pi: number;
+/** @function Element-wise addition, with broadcasting. */
+declare const add: (x: ArrayLike, y: ArrayLike) => Array;
+/** @function Element-wise multiplication, with broadcasting. */
+declare const multiply: (x: ArrayLike, y: ArrayLike) => Array;
+/** @function Numerical negative of every element of an array. */
+declare const negative: (x: ArrayLike) => Array;
+/** @function Calculate element-wise reciprocal of the input. This is `1/x`. */
+declare const reciprocal: (x: ArrayLike) => Array;
+/** @function Round input down to the nearest integer. */
+declare const floor: (x: ArrayLike) => Array;
+/** @function Round input up to the nearest integer. */
+declare const ceil: (x: ArrayLike) => Array;
+/** @function Element-wise sine function (takes radians). */
+declare const sin: (x: ArrayLike) => Array;
+/** @function Element-wise cosine function (takes radians). */
+declare const cos: (x: ArrayLike) => Array;
+/** @function Element-wise inverse sine function (inverse of sin). */
+declare const asin: (x: ArrayLike) => Array;
+/** @function Element-wise inverse tangent function (inverse of tan). */
+declare const atan: (x: ArrayLike) => Array;
+/** @function Calculate the exponential of all elements in the input array. */
+declare const exp: (x: ArrayLike) => Array;
+/** @function Calculate the natural logarithm of all elements in the input array. */
+declare const log: (x: ArrayLike) => Array;
+/** @function Calculate the square root of all elements in the input array. */
+declare const sqrt: (x: ArrayLike) => Array;
+/** @function Return element-wise minimum of the input arrays. */
+declare const minimum: (x: ArrayLike, y: ArrayLike) => Array;
+/** @function Return element-wise maximum of the input arrays. */
+declare const maximum: (x: ArrayLike, y: ArrayLike) => Array;
+/** @function Compare two arrays element-wise. */
+declare const greater: (x: ArrayLike, y: ArrayLike) => Array;
+/** @function Compare two arrays element-wise. */
+declare const less: (x: ArrayLike, y: ArrayLike) => Array;
+/** @function Compare two arrays element-wise. */
+declare const equal: (x: ArrayLike, y: ArrayLike) => Array;
+/** @function Compare two arrays element-wise. */
+declare const notEqual: (x: ArrayLike, y: ArrayLike) => Array;
+/** @function Compare two arrays element-wise. */
+declare const greaterEqual: (x: ArrayLike, y: ArrayLike) => Array;
+/** @function Compare two arrays element-wise. */
+declare const lessEqual: (x: ArrayLike, y: ArrayLike) => Array;
+/** @function Element-wise ternary operator, evaluates to `x` if cond else `y`. */
+declare const where: (cond: ArrayLike, x: ArrayLike, y: ArrayLike) => Array;
 /**
- * Compute the tensor dot product of two N-dimensional arrays.
- *
- * The behavior is determined by `axes`. If an integer `k`, sum over the last
- * `k` axes of x and the first `k` axes of y. If a tuple, then the first array
- * corresponds to the axes of x and the second to the axes of y.
+ * @function
+ * Permute the dimensions of an array. Defaults to reversing the axis order.
  */
-declare function tensordot(x: ArrayLike, y: ArrayLike, axes?: number | [number[], number[]]): Array;
+declare const transpose: (x: ArrayLike, perm?: number[]) => Array;
 /**
- * Einstein summation with string subscripts.
- *
- * @example
- * ```ts
- * import { numpy as np } from "@jax-js/jax";
+ * @function
+ * Give a new shape to an array without changing its data.
  *
- * const a = np.ones([2, 3]);
- * const b = np.ones([3]);
- * np.einsum("ij,j", a, b); // Shape [2]
- * ```
+ * One shape dimension can be -1. In this case, the value is inferred from the
+ * length of the array and remaining dimensions.
+ */
+declare const reshape: (x: ArrayLike, shape: number[]) => Array;
+/**
+ * @function
+ * Move axes of an array to new positions. Other axes retain original order.
  */
-declare function einsum(subscripts: string, ...operands: ArrayLike[]): Array;
+declare const moveaxis: (x: ArrayLike, src: number, dst: number) => Array;
 /**
- * Einstein summation alternating between arrays and numeric indices.
- *
- * @example
- * ```ts
- * import { numpy as np } from "@jax-js/jax";
+ * @function
+ * Add padding (zeros) to an array.
  *
- * const a = np.ones([2, 3]);
- * const b = np.ones([3]);
- * np.einsum(a, [0, 1], b, [1]); // Shape [2]
- * ```
+ * The `width` argument is either an integer or pair of integers, in which case
+ * all axes are padded with the same width. Or if it is an array of pairs, each
+ * pair specifies the padding for its corresponding axis.
  */
-declare function einsum(...args: (ArrayLike | number[])[]): Array;
+declare const pad: (x: ArrayLike, width: number | Pair | Pair[] | Record<number, Pair>) => Array;
 /**
- * Compute the inner product of two arrays.
- *
- * Unlike `jax.numpy.matmul()` or `jax.numpy.dot()`, this always performs a
- * contraction on the last axis.
- *
- * Returned array has shape `[...x.shape[:-1], ...y.shape[:-1]]`.
+ * @function
+ * Return the number of dimensions of an array. Does not consume array reference.
  */
-declare function inner(x: ArrayLike, y: ArrayLike): Array;
+declare const ndim: (x: ArrayLike) => number;
+/** @function Return the shape of an array. Does not consume array reference. */
+declare const shape$1: (x: ArrayLike) => number[];
 /**
- * Compute the outer product of two arrays.
- *
- * If the input arrays are not 1D, they will be flattened. Returned array will
- * be of shape `[x.size, y.size]`.
+ * @function
+ * Return an array of zeros with the same shape and type as a given array.
  */
-declare function outer(x: ArrayLike, y: ArrayLike): Array;
-/** Vector dot product of two arrays along a given axis. */
-declare function vecdot(x: ArrayLike, y: ArrayLike, {
-  axis
-}?: {
-  axis?: number;
-}): Array;
+declare const zerosLike: (a: ArrayLike, dtype?: DType) => Array;
 /**
- * Return the dot product of two vectors.
- *
- * Like vecdot() but flattens the arguments first into vectors.
+ * @function
+ * Return an array of ones with the same shape and type as a given array.
  */
-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;
+declare const onesLike: (a: ArrayLike, dtype?: DType) => Array;
 /**
- * Return a tuple of coordinate matrices from coordinate vectors.
- *
- * Make N-D coordinate arrays for vectorized evaluations of N-D scalar/vector
- * fields over N-D grids, given one-dimensional coordinate arrays x1, x2,…, xn.
+ * @function
+ * Return a full array with the same shape and type as a given array.
  */
-declare function meshgrid(xs: Array[], {
-  indexing
-}?: {
-  indexing?: "xy" | "ij";
-}): Array[];
+declare const fullLike: (a: ArrayLike, fillValue: number | boolean | Array, dtype?: DType) => Array;
 /**
- * Clip (limit) the values in an array.
- *
- * Given an interval, values outside the interval are clipped to the interval
- * edges. For example, if an interval of [0, 1] is specified, values smaller
- * than 0 become 0, and values larger than 1 become 1.
- *
- * If either bound is undefined, it is ignored.
+ * Return the number of elements in an array, optionally along an axis.
+ * Does not consume array reference.
  */
-declare function clip(a: ArrayLike, min?: ArrayLike, max?: ArrayLike): Array;
+declare function size(a: ArrayLike, axis?: number): number;
+/** Convert an array to a specified dtype. */
+declare function astype(a: ArrayLike, dtype: DType): Array;
+/** Sum of the elements of the array over a given axis, or axes. */
+declare function sum(a: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
+/** Product of the array elements over a given axis. */
+declare function prod(a: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
+/** Return the minimum of array elements along a given axis. */
+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;
 /**
- * Calculate the absolute value element-wise.
+ * Test whether any array element along a given axis evaluates to True.
  *
- * This is the same function as `jax.numpy.abs()`.
+ * Returns a boolean array with the same shape as `a` with the specified axis
+ * removed. If axis is None, returns a scalar.
  */
-declare function absolute(x: ArrayLike): Array;
-/** Return an element-wise indication of sign of the input. */
-declare function sign(x: ArrayLike): Array;
-/** @function Return element-wise positive values of the input (no-op). */
-declare const positive: (x: ArrayLike) => Array;
+declare function any(a: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
 /**
- * Return the Hamming window of size M, a taper with a weighted cosine bell.
+ * Test whether all array elements along a given axis evaluate to True.
  *
- * `w(n) = 0.54 - 0.46 * cos(2πn/(M-1))` for `0 <= n <= M-1`.
+ * Returns a boolean array with the same shape as `a` with the specified axis
+ * removed. If axis is None, returns a scalar.
  */
-declare function hamming(M: number): Array;
+declare function all(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. */
+declare function mean(a: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
 /**
- * Return the Hann window of size M, a taper with a weighted cosine bell.
+ * Returns the indices of the minimum values along an axis.
  *
- * `w(n) = 0.5 - 0.5 * cos(2πn/(M-1))` for `0 <= n <= M-1`.
+ * By default, index is into the flatted array, otherwise it is along the
+ * specified axis.
  */
-declare function hann(M: number): Array;
+declare function argmin(a: ArrayLike, axis?: number, opts?: ReduceOpts): Array;
 /**
- * @function
- * Compute the Heaviside step function. It is defined piecewise:
- * - `heaviside(x1, x2) = 0` for `x1 < 0`,
- * - `heaviside(x1, x2) = x2` for `x1 == 0`,
- * - `heaviside(x1, x2) = 1` for `x1 > 0`.
+ * Returns the indices of the maximum values along an axis.
+ *
+ * By default, index is into the flatted array, otherwise it is along the
+ * specified axis.
  */
-declare const heaviside: OwnedFunction<(x1: ArrayLike, x2: ArrayLike) => Array>;
-/** Calculate element-wise square of the input array. */
-declare function square(x: ArrayLike): Array;
-/** Element-wise tangent function (takes radians). */
-declare function tan(x: ArrayLike): Array;
+declare function argmax(a: ArrayLike, axis?: number, opts?: ReduceOpts): 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.
+ * Cumulative sum of elements along an axis.
  *
- * **Note:** JVP is not supported at x=0 due to discontinuous derivative. This
- * requires a custom JVP rule to handle properly (see JAX implementation).
+ * Currently this function is `O(n^2)`, we'll improve this later on with a
+ * two-phase parallel reduction algorithm.
  */
-declare const sinc: OwnedFunction<(x: ArrayLike) => Array>;
-/** Element-wise inverse cosine function (inverse of cos). */
-declare function acos(x: ArrayLike): 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;
 /**
- * @function
- * Return element-wise hypotenuse for the given legs of a right triangle.
+ * Split an array into multiple sub-arrays along an axis.
  *
- * In the original NumPy/JAX implementation, this function is more numerically
- * stable than `sqrt(x1**2 + x2**2)`. We don't currently implement those
- * stability improvements.
+ * @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 const hypot: OwnedFunction<(x1: ArrayLike, x2: ArrayLike) => Array>;
+declare function split$1(a: ArrayLike, indicesOrSections: number | number[], axis?: number): Array[];
 /**
- * @function
- * Element-wise arc tangent of y/x with correct quadrant.
- *
- * Returns the angle in radians between the positive x-axis and the point (x, y).
- * The result is in the range [-π, π].
+ * Join a sequence of arrays along an existing axis.
  *
- * Uses numerically stable formulas:
- * - When x >= 0: atan2(y, x) = 2 * atan(y / (sqrt(x^2 + y^2) + x))
- * - When x < 0:  atan2(y, x) = 2 * atan((sqrt(x^2 + y^2) - x) / y)
+ * The arrays must have the same shape, except in the dimension corresponding to
+ * `axis` (the first, by default).
  *
- * The output is ill-defined when both x and y are zero.
+ * No scalars can be passed to this function, as the axis is then ambiguous.
  */
-declare const atan2: 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;
+declare function concatenate(xs: Array[], axis?: number): Array;
 /**
- * Return the largest integer smaller or equal to the division of the inputs.
- *
- * The result is always rounded towards negative infinity.
+ * Join a sequence of arrays along a new axis.
  *
- * 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).
+ * The `axis` parameter specifies the index of the new axis in the dimensions of
+ * the result. For example, if `axis=0` it will be the first dimension and if
+ * `axis=-1` it will be the last dimension.
  *
- * @param x - Dividend array.
- * @param y - Divisor array.
- * @returns Element-wise floor division of x by y.
+ * All shapes must have the same shape.
  */
-declare function floorDivide(x: ArrayLike, y: ArrayLike): Array;
+declare function stack(xs: ArrayLike[], axis?: number): Array;
 /**
- * @function
- * Calculate element-wise floating-point modulo operation.
+ * Horizontally stack arrays. Inputs are promoted to rank at least 1, then
+ * concatenated along axis 1 (if rank-2 or higher) or 0 (if rank-1).
  */
-declare const fmod: OwnedFunction<(x: ArrayLike, y: ArrayLike) => Array>;
+declare function hstack(xs: ArrayLike[]): Array;
 /**
- * @function
- * Calculate element-wise remainder of the division (matches sign of y).
+ * Vertically stack arrays. Inputs are promoted to rank at least 2, then
+ * concatenated along axis 0.
  */
-declare const remainder: OwnedFunction<(x: ArrayLike, y: ArrayLike) => Array>;
+declare function vstack(xs: ArrayLike[]): Array;
 /**
- * Return element-wise quotient and remainder simultaneously.
+ * Stack arrays depth-wise. Inputs are promoted to rank at least 3, then
+ * concatenated along axis 2.
+ */
+declare function dstack(xs: ArrayLike[]): Array;
+/**
+ * Stack arrays column-wise. Inputs are promoted to rank at least 2, then
+ * concatenated along axis 1.
+ */
+declare function columnStack(xs: ArrayLike[]): Array;
+/** Flip an array vertically (axis=0). */
+declare function flipud(x: ArrayLike): Array;
+/** Flip an array horizontally (axis=1). */
+declare function fliplr(x: ArrayLike): Array;
+/** Interchange two axes of an array. */
+declare function swapaxes(a: ArrayLike, axis1: number, axis2: number): Array;
+/** Transpose the last two dimensions of an array. */
+declare function matrixTranspose(a: ArrayLike): Array;
+/** Return a 1-D flattened array containing the elements of the input. */
+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.
  *
- * Equivalent to `[floorDivide(x, y), remainder(x, y)]`.
+ * @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.
  *
- * @param x - Dividend array.
- * @param y - Divisor array.
- * @returns Tuple of [quotient, remainder].
+ * @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 divmod(x: ArrayLike, y: ArrayLike): [Array, Array];
-/** Round input to the nearest integer towards zero. */
-declare function trunc(x: ArrayLike): Array;
+declare function expandDims(a: ArrayLike, axis: number | number[]): Array;
+/**
+ * Repeat each element of an array after themselves.
+ *
+ * If no axis is provided, use the flattened input array, and return a flat
+ * output array.
+ */
+declare function repeat(a: ArrayLike, repeats: number, axis?: number): Array;
 /**
- * Compute `x1 * 2 ** x2` as a standard multiplication and exponentiation.
+ * Construct an array by repeating A the number of times given by reps.
  *
- * This is the inverse of `frexp()`.
+ * If `A` is an array of shape `(d1, d2, ..., dn)` and `reps` is a sequence of
+ * integers, the resulting array will have a shape of `(reps[0] * d1,
+ * reps[1] * d2, ..., reps[n] * dn)`, with `A` tiled along each dimension.
  */
-declare function ldexp(x1: ArrayLike, x2: ArrayLike): Array;
+declare function tile(a: ArrayLike, reps: number | number[]): Array;
 /**
- * Decompose floating-point values into mantissa and two's exponent.
+ * Broadcast an array to a shape, with NumPy-style broadcasing rules.
  *
- * The mantissa is returned in the range `(-1, 1)` with magnitude `>= 0.5` if
- * `x != 0`, and the exponent is an integer such that
- * `x = mantissa * 2**exponent`.
+ * In other words, this lets you append axes to the left, and/or expand
+ * dimensions where the shape is 1.
  */
-declare function frexp(x: ArrayLike): [Array, Array];
-/** Calculate `2**p` for all p in the input array. */
-declare function exp2(p: ArrayLike): Array;
-/** Return the base-2 logarithm of x, element-wise. */
-declare function log2(x: ArrayLike): Array;
-/** Return the base-10 logarithm of x, element-wise. */
-declare function log10(x: ArrayLike): Array;
-/** Calculate `exp(x) - 1` element-wise. */
-declare function expm1(x: ArrayLike): Array;
-/** Calculate the natural logarithm of `1 + x` element-wise. */
-declare function log1p(x: ArrayLike): Array;
-/** Convert angles from degrees to radians. */
-declare function deg2rad(x: ArrayLike): Array;
-/** @function Alias of `jax.numpy.deg2rad()`. */
-declare const radians: typeof deg2rad;
-/** Convert angles from radians to degrees. */
-declare function rad2deg(x: ArrayLike): Array;
-/** @function Alias of `jax.numpy.rad2deg()`. */
-declare const degrees: typeof rad2deg;
+declare function broadcastTo(a: ArrayLike, shape: number[]): Array;
+/** Broadcast input shapes to a common output shape. */
+declare function broadcastShapes(...shapes: number[][]): number[];
+/** Broadcast arrays to a common shape. */
+declare function broadcastArrays(...arrays: ArrayLike[]): Array[];
 /**
- * @function
- * Computes first array raised to power of second array, element-wise.
+ * Return specified diagonals.
+ *
+ * If a is 2D, return the diagonal of the array with the given offset. If a is
+ * 3D or higher, compute diagonals along the two given axes (default: 0, 1).
+ *
+ * This returns a view over the existing array. The shape of the resulting array
+ * is determined by removing the two axes along which the diagonal is taken,
+ * then appending a new axis to the right with holding the diagonals.
  */
-declare const power: OwnedFunction<(x1: ArrayLike, x2: ArrayLike) => Array>;
-/** @function Calculate the element-wise cube root of the input array. */
-declare const cbrt: OwnedFunction<(x: ArrayLike) => Array>;
+declare function diagonal(a: ArrayLike, offset?: number, axis1?: number, axis2?: number): Array;
 /**
- * @function
- * Calculate element-wise hyperbolic sine of input.
+ * Extract a diagonal or construct a diagonal array.
  *
- * `sinh(x) = (exp(x) - exp(-x)) / 2`
+ * If v is a 2D array, return the k-th diagonal of v (as a view). If v is a 1D
+ * array, return a 2D array with v on the k-th diagonal.
  */
-declare const sinh: OwnedFunction<(x: ArrayLike) => Array>;
+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;
 /**
- * @function
- * Calculate element-wise hyperbolic cosine of input.
+ * Return a sorted copy of an array.
  *
- * `cosh(x) = (exp(x) + exp(-x)) / 2`
+ * 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 const cosh: OwnedFunction<(x: ArrayLike) => Array>;
+declare function sort(a: ArrayLike, axis?: number): Array;
 /**
- * @function
- * Calculate element-wise hyperbolic tangent of input.
+ * Return indices that would sort an array. This may be an unstable sorting
+ * algorithm; it need not preserve order of indices in ties.
  *
- * `tanh(x) = sinh(x)/cosh(x) = (exp(x) - exp(-x)) / (exp(x) + exp(-x))`
+ * Returns an array of `int32` indices.
+ *
+ * The array is sorted along a specified axis (the last by default).
  */
-declare const tanh: OwnedFunction<(x: ArrayLike) => Array>;
+declare function argsort(a: ArrayLike, axis?: number): Array;
 /**
- * @function
- * Calculate element-wise inverse hyperbolic sine of input.
+ * Take elements from an array along an axis.
  *
- * `arcsinh(x) = ln(x + sqrt(x^2 + 1))`
+ * This is equivalent to advanced indexing with integer indices over that
+ * numbered axis. By default, the flattened array is used.
  */
-declare const arcsinh: OwnedFunction<(x: ArrayLike) => Array>;
+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;
+  atol?: number;
+}): boolean;
+/** Matrix product of two arrays. */
+declare function matmul(x: ArrayLike, y: ArrayLike): Array;
+/** Dot product of two arrays. */
+declare function dot(x: ArrayLike, y: ArrayLike): Array;
 /**
- * @function
- * Calculate element-wise inverse hyperbolic cosine of input.
+ * Compute the tensor dot product of two N-dimensional arrays.
  *
- * `arccosh(x) = ln(x + sqrt(x^2 - 1))`
+ * The behavior is determined by `axes`. If an integer `k`, sum over the last
+ * `k` axes of x and the first `k` axes of y. If a tuple, then the first array
+ * corresponds to the axes of x and the second to the axes of y.
  */
-declare const arccosh: OwnedFunction<(x: ArrayLike) => Array>;
+declare function tensordot(x: ArrayLike, y: ArrayLike, axes?: number | [number[], number[]]): Array;
 /**
- * @function
- * Calculate element-wise inverse hyperbolic tangent of input.
+ * Einstein summation with string subscripts.
  *
- * `arctanh(x) = 0.5 * ln((1 + x) / (1 - x))`
+ * @example
+ * ```ts
+ * import { numpy as np } from "@jax-js/jax";
+ *
+ * const a = np.ones([2, 3]);
+ * const b = np.ones([3]);
+ * np.einsum("ij,j", a, b); // Shape [2]
+ * ```
  */
-declare const arctanh: OwnedFunction<(x: ArrayLike) => Array>;
+declare function einsum(subscripts: string, ...operands: ArrayLike[]): Array;
 /**
- * Compute the variance of an array.
+ * Einstein summation alternating between arrays and numeric indices.
  *
- * The variance is computed for the flattened array by default, otherwise over
- * the specified axis.
+ * @example
+ * ```ts
+ * import { numpy as np } from "@jax-js/jax";
  *
- * If `correction` is provided, the divisor in calculation is `N - correction`,
- * where `N` represents the number of elements (e.g., for Bessel's correction).
+ * const a = np.ones([2, 3]);
+ * const b = np.ones([3]);
+ * np.einsum(a, [0, 1], b, [1]); // Shape [2]
+ * ```
  */
-declare function var_(x: ArrayLike, axis?: Axis, opts?: {
-  mean?: ArrayLike;
-  correction?: number;
-} & ReduceOpts): Array;
+declare function einsum(...args: (ArrayLike | number[])[]): Array;
 /**
- * Compute the standard deviation of an array.
+ * Compute the inner product of two arrays.
  *
- * The standard deviation is computed for the flattened array by default,
- * otherwise over the specified axis.
+ * Unlike `jax.numpy.matmul()` or `jax.numpy.dot()`, this always performs a
+ * contraction on the last axis.
  *
- * If `correction` is provided, the divisor in calculation is `N - correction`,
- * where `N` represents the number of elements (e.g., for Bessel's correction).
+ * Returned array has shape `[...x.shape[:-1], ...y.shape[:-1]]`.
  */
-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
+declare function inner(x: ArrayLike, y: ArrayLike): Array;
+/**
+ * Compute the outer product of two arrays.
+ *
+ * If the input arrays are not 1D, they will be flattened. Returned array will
+ * be of shape `[x.size, y.size]`.
+ */
+declare function outer(x: ArrayLike, y: ArrayLike): Array;
+/** Vector dot product of two arrays along a given axis. */
+declare function vecdot(x: ArrayLike, y: ArrayLike, {
+  axis
 }?: {
-  rowvar?: boolean;
+  axis?: number;
 }): 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). */
-declare function isnan(x: ArrayLike): Array;
-/** Test element-wise for negative infinity, return bool array. */
-declare function isneginf(x: ArrayLike): Array;
-/** Test element-wise for positive infinity, return bool array. */
-declare function isposinf(x: ArrayLike): Array;
 /**
- * @function
- * Test element-wise for finite values (not infinity or NaN).
+ * Return the dot product of two vectors.
+ *
+ * Like vecdot() but flattens the arguments first into vectors.
  */
-declare const isfinite: OwnedFunction<(x: ArrayLike) => Array>;
-declare namespace tree_d_exports {
-  export { JsTree, JsTreeDef, MapJsTree, NodeType, dispose, flatten, leaves, map, ref, structure, unflatten };
-}
-declare enum NodeType {
-  Array = "Array",
-  Object = "Object",
-  Leaf = "Leaf",
-}
-/** Analog to the JAX "pytree" object, but for JavaScript. */
-type JsTree<T> = T | JsTree<T>[] | {
-  [key: string]: JsTree<T>;
-};
-type Same<X, Y> = (<T>() => T extends X ? 1 : 2) extends (<T>() => T extends Y ? 1 : 2) ? true : false;
-type MappedJsTree<T, A, B> = T extends A ? B : T extends Array ? T : T extends globalThis.Array<infer U> ? number extends T["length"] ? MapJsTree<U, A, B>[] : { [K in keyof T]: MapJsTree<T[K], A, B> } : { [K in keyof T]: MapJsTree<T[K], A, B> };
-/** @ignore Convert a subtype of JsTree<A> into a JsTree<B>, with the same structure. */
-type MapJsTree<T, A, B> = Same<A, B> extends true ? T : MappedJsTree<T, A, B>;
-/** Represents the structure of a JsTree. */
-declare class JsTreeDef {
-  readonly nodeType: NodeType;
-  readonly nodeMetadata: any;
-  readonly childTreedefs: JsTreeDef[];
-  static leaf: JsTreeDef;
-  constructor(nodeType: NodeType, nodeMetadata: any,
-  // Must be comparable with deepEqual.
-  childTreedefs: JsTreeDef[]);
-  /** Get the total number of leaves in the tree. */
-  get size(): number;
-  /** Returns a string representation of this tree definition. */
-  toString(root?: boolean): string;
-  /** Compare this tree definition with another. */
-  equals(other: JsTreeDef): boolean;
-}
-/** Flatten a structured object, returning the tree definition. */
-declare function flatten<T>(tree: JsTree<T>): [T[], JsTreeDef];
-/** Get the leaves of a tree. */
-declare function leaves<T>(tree: JsTree<T>): T[];
-/** Get the treedef for a tree. */
-declare function structure<T>(tree: JsTree<T>): JsTreeDef;
-/** Reconstruct a structured object from the flattened representation. */
-declare function unflatten<T>(treedef: JsTreeDef, leaves: Iterable<T>): JsTree<T>;
-/** Maps a multi-input function over pytree args to produce a new pytree. */
-declare function map<T, U, Tree extends JsTree<T>>(fn: (...args: T[]) => U, tree: Tree, ...rest: Tree[]): MapJsTree<Tree, T, U>;
-/** Take a reference of every array in a tree. */
-declare function ref<Tree extends JsTree<Array>>(tree: Tree): Tree;
-/** Dispose every array in a tree. */
-declare function dispose<Tree extends JsTree<Array>>(tree: Tree | null | undefined): void;
-//#endregion
-//#region src/frontend/convolution.d.ts
-/** Definition of a general dilated convolution. Should be valid on creation. */
-interface ConvParams {
-  vmapDims: number;
-  strides: number[];
-  padding: Pair[];
-  lhsDilation: number[];
-  rhsDilation: number[];
-}
+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.
+ *
+ * Make N-D coordinate arrays for vectorized evaluations of N-D scalar/vector
+ * fields over N-D grids, given one-dimensional coordinate arrays x1, x2,…, xn.
+ */
+declare function meshgrid(xs: Array[], {
+  indexing
+}?: {
+  indexing?: "xy" | "ij";
+}): Array[];
 /**
- * Check that the shapes and parameters passed to convolution are valid.
- * Expected shapes of the lhs and rhs of the convolution are:
+ * Clip (limit) the values in an array.
  *
- * - `lhsShape = [*vmapDims, batchSize, inChannels, spatialDims...]`
- * - `rhsShape = [*vmapDims, outChannels, inChannels, kernelSize...]`
+ * Given an interval, values outside the interval are clipped to the interval
+ * edges. For example, if an interval of [0, 1] is specified, values smaller
+ * than 0 become 0, and values larger than 1 become 1.
  *
- * If the check succeeds, returns the output shape.
+ * If either bound is undefined, it is ignored.
  */
-//#endregion
-//#region src/frontend/jaxpr.d.ts
+declare function clip(a: ArrayLike, min?: ArrayLike, max?: ArrayLike): Array;
 /**
- * Function callback with an associated dispose() method.
+ * Calculate the absolute value element-wise.
  *
- * The dispose() method should be called to clean up any tracer resources needed
- * by the function after the last time it is called.
+ * This is the same function as `jax.numpy.abs()`.
  */
-type OwnedFunction<F extends Function> = F & {
-  dispose: () => void;
-};
-/** Variable in a Jaxpr expression. */
-declare class Var {
-  #private;
-  readonly id: number;
-  readonly aval: ShapedArray;
-  constructor(aval: ShapedArray);
-  toString(): string;
-}
-/** Literal in a Jaxpr expression. Currently, only scalars are supported. */
-declare class Lit {
-  readonly value: number;
-  readonly aval: ShapedArray;
-  get dtype(): DType;
-  constructor(aval: AbstractValue, value: number);
-}
-type Atom = Var | Lit;
-declare class VarPrinter {
-  #private;
-  names: Map<Var, string>;
-  name(v: Var): string;
-  nameType(v: Var): string;
-}
-/** A single statement / binding in a Jaxpr, in ANF form. */
-declare class JaxprEqn {
-  readonly primitive: Primitive;
-  readonly inputs: Atom[];
-  readonly params: Record<string, any>;
-  readonly outBinders: Var[];
-  constructor(primitive: Primitive, inputs: Atom[], params: Record<string, any>, outBinders: Var[]);
-  pprint(usedVars?: Set<Var>, vp?: VarPrinter): PPrint;
-  toString(): string;
-}
-/** Typed intermediate representation for traced computations. */
-declare class Jaxpr implements FpHashable {
-  #private;
-  readonly inBinders: Var[];
-  readonly eqns: JaxprEqn[];
-  readonly outs: Atom[];
-  constructor(inBinders: Var[], eqns: JaxprEqn[], outs: Atom[]);
-  pprint(): PPrint;
-  toString(): string;
-  /**
-   * Gets a hash of this Jaxpr.
-   *
-   * Var identity is not considered in the hash, so two Jaxprs with the same
-   * order of assignments and operators but different variable IDs will resolve
-   * to the same hash (and toString representation).
-   */
-  getHash(): bigint;
-  hash(state: FpHash): void;
-  /**
-   * Produce a simplified Jaxpr with basic optimizations applied.
-   *  - Trim away unused variables.
-   *  - Fold away *1, *0, or +0 operations against literals.
-   *  - Remove no-op movement operations.
-   */
-  simplify(): Jaxpr;
-  /** 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[];
-};
-//#endregion
-//#region src/frontend/core.d.ts
+declare function absolute(x: ArrayLike): Array;
+/** Return an element-wise indication of sign of the input. */
+declare function sign(x: ArrayLike): Array;
+/** @function Return element-wise positive values of the input (no-op). */
+declare const positive: (x: ArrayLike) => Array;
 /**
- * Frontend primitive operations, which are lowered into Kernel objects before
- * being dispatched to the backend.
- *
- * Any operation between arrays can be described in these parts. This is also
- * the set of primitives that can occur in Jaxpr programs, and the level at
- * which transformations like vmap, grad, and jvp occur. They are loosely based
- * on [XLA](https://openxla.org/xla/operation_semantics).
- *
- * All n-ary operations support broadcasting, with NumPy semantics.
- */
-declare enum Primitive {
-  Add = "add",
-  Mul = "mul",
-  Idiv = "idiv",
-  Mod = "mod",
-  // uses sign of numerator, C-style, matches JS but not Python
-  Min = "min",
-  Max = "max",
-  Neg = "neg",
-  Reciprocal = "reciprocal",
-  Floor = "floor",
-  Ceil = "ceil",
-  StopGradient = "stop_gradient",
-  Cast = "cast",
-  Bitcast = "bitcast",
-  Sin = "sin",
-  Cos = "cos",
-  Asin = "asin",
-  Atan = "atan",
-  Exp = "exp",
-  Log = "log",
-  Erf = "erf",
-  Erfc = "erfc",
-  Sqrt = "sqrt",
-  Reduce = "reduce",
-  Dot = "dot",
-  // sum(x*y, axis=-1)
-  Conv = "conv",
-  // see lax.conv_general_dilated
-  Pool = "pool",
-  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",
-  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]: {
-    dtype: DType;
-  };
-  [Primitive.Bitcast]: {
-    dtype: DType;
-  };
-  [Primitive.Reduce]: {
-    op: AluOp;
-    axis: number[];
-  };
-  [Primitive.Conv]: ConvParams;
-  [Primitive.Pool]: {
-    window: number[];
-    strides: number[];
-  };
-  [Primitive.PoolTranspose]: {
-    inShape: number[];
-    window: number[];
-    strides: number[];
-  };
-  [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[];
-  };
-  [Primitive.Broadcast]: {
-    shape: number[];
-    axis: number[];
-  };
-  [Primitive.Reshape]: {
-    shape: number[];
-  };
-  [Primitive.Flip]: {
-    axis: number[];
-  };
-  [Primitive.Shrink]: {
-    slice: Pair[];
-  };
-  [Primitive.Pad]: {
-    width: Pair[];
-  };
-  [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>;
-declare enum CompareOp {
-  Less = "less",
-  Equal = "equal",
-  NotEqual = "not_equal",
-  LessEqual = "less_equal",
-}
-/** @inline */
-type Axis = number | number[] | null;
-/** @inline */
-type ReduceOpts = {
-  keepdims?: boolean;
-};
-type MainTrace = {
-  level: number;
-  traceType: new (main: MainTrace) => Trace;
-  globalData: any | null;
-};
+ * Return the Hamming window of size M, a taper with a weighted cosine bell.
+ *
+ * `w(n) = 0.54 - 0.46 * cos(2πn/(M-1))` for `0 <= n <= M-1`.
+ */
+declare function hamming(M: number): Array;
 /**
- * Push an interpreter onto the trace stack. Use this like:
- * `using main = newMain(...);`
+ * Return the Hann window of size M, a taper with a weighted cosine bell.
+ *
+ * `w(n) = 0.5 - 0.5 * cos(2πn/(M-1))` for `0 <= n <= M-1`.
  */
-
-type TracerValue = Tracer | number | boolean;
-declare abstract class Trace {
-  readonly main: MainTrace;
-  constructor(main: MainTrace);
-  abstract pure(val: TracerValue): Tracer;
-  abstract lift(val: Tracer): Tracer;
-  abstract processPrimitive<P extends Primitive>(primitive: P, tracers: Tracer[], params: PrimitiveParams<P>): Tracer[];
-}
-/** Internal representation of an array value. */
-interface AbstractValue {
-  /** Shape of the array. Must be a static tuple of non-negative dimensions. */
-  shape: number[];
-  /** Concrete data type of array elements. */
-  dtype: DType;
-  /**
-   * Arrays created from JavaScript numbers (e.g., `np.array(3)`) are created as
-   * _weakly typed_ unless a dtype is explicitly specified.
-   *
-   * Weakly typed values will automatically cast to the data type of other
-   * arrays when used as an operand as an expression. This property only affects
-   * how they promote in type casting; their memory layout is still determined
-   * by the actual `dtype` field.
-   *
-   * ```ts
-   * const x = np.array(3); // weakType = true, dtype = float32
-   * const y = np.array([1, 2], { dtype: np.int32 }); // weakType = false, dtype = int32
-   * const z = x.add(y); // z has dtype int32 because x is weakly typed
-   * ```
-   *
-   * Weak types are present in JIT programs in their spec (e.g., Jaxpr inputs
-   * and outputs can be weakly typed) form. But they're solely a frontend
-   * concept. Backends are not aware of weak types.
-   */
-  weakType: boolean;
-}
+declare function hann(M: number): Array;
 /**
- * Broadcast shapes and promote types with casting for two avals.
+ * @function
+ * Compute the Heaviside step function. It is defined piecewise:
+ * - `heaviside(x1, x2) = 0` for `x1 < 0`,
+ * - `heaviside(x1, x2) = x2` for `x1 == 0`,
+ * - `heaviside(x1, x2) = 1` for `x1 > 0`.
+ */
+declare const heaviside: OwnedFunction<(x1: ArrayLike, x2: ArrayLike) => Array>;
+/** Calculate element-wise square of the input 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.
  *
- * This implements the weak type behavior described in `promoteTypes()`, but not
- * implemented in that function as `weakType` is not passed.
+ * 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 abstract class Tracer {
-  /** @ignore */
-  readonly _trace: Trace;
-  constructor(trace: Trace);
-  abstract get aval(): AbstractValue;
-  abstract toString(): string;
-  /**
-   * Access an array by reference, incrementing the reference count.
-   *
-   * jax-js handles freeing arrays by using "move" semantics, like in Rust/C++.
-   * Whenever you pass an array into a function, that function should consume
-   * the array, and it will no longer be usable. For example, if you had:
-   *
-   * ```
-   * const x = np.array([1, 2, 3]);
-   * const y = np.add(x, x);
-   * ```
-   *
-   * The second line does not work because the first parameter consumes `x`, and
-   * then the second parameter will already have been freed / disposed.
-   *
-   * To fix this, you can write:
-   *
-   * ```
-   * const y = np.add(x.ref, x);
-   * ```
-   *
-   * Under the hood, every access to `.ref` increments the internal reference
-   * count of the array. The reference count starts at 1. When it hits 0, the
-   * memory behind the array is freed.
-   */
-  abstract get ref(): this;
-  /**
-   * Manually decrement the reference count of the array.
-   *
-   * Arrays are created with reference count 1. Whenever it is used as argument
-   * to a function or other operation, it is disposed (i.e., reference count
-   * decreases by 1) automatically. Whenever a `.ref` is created, the reference
-   * count increases.
-   *
-   * You generally don't need to call this function directly since arrays are
-   * automatically disposed after being passed into an operation. One common
-   * exception is when writing a function and ignoring one of its arguments. In
-   * that case, by convention you should dispose of that argument manually.
-   *
-   * ```
-   * function myCustomOperation(a: np.Array, b: np.Array) {
-   *   b.dispose(); // Needed to satisfy "move" rules.
-   *   return a.add(1);
-   * }
-   * ```
-   */
-  abstract dispose(): void;
-  /** The shape of the array. */
-  get shape(): number[];
-  /** The total number of elements in the array. */
-  get size(): number;
-  /** The dtype of elements stored in the array. */
-  get dtype(): DType;
-  /**
-   * Whether the array is weakly typed.
-   *
-   * Weakly typed arrays will cast to the dtype of the other operand. See
-   * `promoteTypes()` for details.
-   */
-  get weakType(): boolean;
-  /** The number of dimensions of the array. */
-  get ndim(): number;
-  /** @ignore */
-  fullLower(): 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;
-  notEqual(other: this | TracerValue): this;
-  greaterEqual(other: this | TracerValue): this;
-  lessEqual(other: this | TracerValue): this;
-  /** Sum of the elements of the array over a given axis, or axes. */
-  sum(axis?: Axis, opts?: ReduceOpts): this;
-  /** Product of the array elements over a given axis. */
-  prod(axis?: Axis, opts?: ReduceOpts): this;
-  /** Compute the average of the array elements along the specified axis. */
-  mean(axis?: Axis, opts?: ReduceOpts): this;
-  /** Permute the dimensions of an array. Defaults to reversing the axis order. */
-  transpose(perm?: number[]): this;
-  /**
-   * Give a new shape to an array without changing its data.
-   *
-   * One shape dimension can be -1. In this case, the value is inferred from the
-   * length of the array and remaining dimensions.
-   */
-  reshape(shape: number | number[]): this;
-  /** Copy the array and cast to a specified dtype. */
-  astype(dtype: DType): this;
-  /** Subtract an array from this one. */
-  sub(other: this | TracerValue): this;
-  /** Divide an array by this one. */
-  div(other: this | TracerValue): this;
-  /** 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;
-  /** Flatten the array without changing its data. */
-  ravel(): this;
-  /**
-   * Iterate over the first dimension of this array, returning slices.
-   *
-   * This can be used to destructure arrays. For example:
-   *
-   * ```js
-   * let x = np.array([[1, 2], [3, 4]]);
-   * let [a, b] = x;
-   * console.log(a.js()); // [1, 2]
-   * console.log(b.js()); // [3, 4]
-   * ```
-   */
-  [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.
-   *
-   * This is the equivalent of slicing in Python, e.g. `x[1:3, 2, :, None]`. To
-   * mimic this in JavaScript, we would write:
-   *
-   * ```js
-   * x.slice([1, 3], 2, [], null);
-   * ```
-   *
-   * The `slice` method accepts a variable number of arguments, each of which
-   * can be a number, an empty array, a single-element array, a two-element
-   * array, or `null`. The arguments are interpreted as follows:
-   *
-   * - A number `n` means to access the `n`-th element along that axis, removing
-   *   that axis from the resulting shape.
-   * - An empty array `[]` means to keep that axis as-is, like `:` in Python.
-   * - A single-element array `[i]` means to start slicing from index `i`
-   *   (inclusive) to the end of the axis, like `x[i:]`.
-   * - A two-element array `[i, j]` means to slice from index `i` (inclusive)
-   *   to index `j` (exclusive), like `x[i:j]`.
-   * - `null` means to add a new axis at that position, like `np.newaxis`.
-   *
-   * Like in Python, negative indices are supported, which count from the end of
-   * the axis. For example, `-1` means the last element.
-   *
-   * Strided slices are not yet implemented, so you cannot write `x[::2]` or
-   * similar.
-   *
-   * Advanced indexing by integer arrays is also supported. This translates to
-   * the "gather" primitive, and it allows you to access specific elements of
-   * the array by integer indices stored in another array.
-   */
-  slice(...index: (number | [] | [number] | Pair | null | Tracer)[]): this;
-}
-declare class ShapedArray implements AbstractValue {
-  readonly shape: number[];
-  readonly dtype: DType;
-  readonly weakType: boolean;
-  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;
-}
-//#endregion
-//#region src/frontend/array.d.ts
-type ArrayLike = Array | number | boolean;
-/** Version of pureArray with fudged types. */
-
+declare const sinc: OwnedFunction<(x: ArrayLike) => Array>;
+/** Element-wise inverse cosine function (inverse of cos). */
+declare function acos(x: ArrayLike): Array;
 /**
- * An executable operation that will be dispatched to the backend.
+ * @function
+ * Return element-wise hypotenuse for the given legs of a right triangle.
  *
- * This holds a reference to all input buffers used in the operation. After the
- * operation is dispatched, the references should be released.
+ * In the original NumPy/JAX implementation, this function is more numerically
+ * stable than `sqrt(x1**2 + x2**2)`. We don't currently implement those
+ * stability improvements.
  */
-declare class PendingExecute {
-  #private;
-  readonly backend: Backend;
-  readonly source: Kernel | Routine;
-  readonly inputs: Slot[];
-  readonly outputs: Slot[];
-  prepared: Executable | null;
-  submitted: boolean;
-  constructor(backend: Backend, source: Kernel | Routine, inputs: Slot[], outputs: Slot[]);
-  updateRc(delta: number): void;
-  prepare(): Promise<void>;
-  prepareSync(): void;
-  submit(): void;
-}
-/** @inline */
-type DTypeAndDevice = {
-  dtype?: DType;
-  device?: Device;
-};
-type ArrayConstructorArgs = {
-  source: AluExp | Slot;
-  st: ShapeTracker;
-  dtype: DType;
-  weakType: boolean;
-  backend: Backend;
-  committed: boolean;
-  pending?: Iterable<PendingExecute>;
-};
+declare const hypot: OwnedFunction<(x1: ArrayLike, x2: ArrayLike) => Array>;
 /**
- * A multidimensional numeric array with data stored on CPU or GPU.
+ * @function
+ * Element-wise arc tangent of y/x with correct quadrant.
  *
- * This is the library's core data type. Equivalent to `jax.Array` from JAX, or
- * `torch.Tensor`.
+ * Returns the angle in radians between the positive x-axis and the point (x, y).
+ * The result is in the range [-π, π].
  *
- * Not to be confused with the JavaScript "Array" constructor. Avoid importing
- * this into your code's namespace if you're already using the JavaScript
- * "Array" type by name.
+ * Uses numerically stable formulas:
+ * - When x >= 0: atan2(y, x) = 2 * atan(y / (sqrt(x^2 + y^2) + x))
+ * - When x < 0:  atan2(y, x) = 2 * atan((sqrt(x^2 + y^2) - x) / y)
+ *
+ * The output is ill-defined when both x and y are zero.
  */
-declare class Array extends Tracer {
-  #private;
-  /**
-   * @ignore
-   * Constructs an array from source, shape and backend. Note that if the source
-   * is a backend `Slot`, this constructor _takes ownership_ of the slot. It
-   * will be freed when the array is disposed.
-   */
-  constructor(args: ArrayConstructorArgs);
-  /** @ignore */
-  get aval(): ShapedArray;
-  /** Return a simple string representation of the array's dimensions. */
-  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.
-   *
-   * This only works for scalars (0-dimensional arrays). It lets you get values
-   * "out" of the JAX system. For instance, if `x = np.array(5)`, then you can
-   * evaluate `x + 1` and `x ** 2` to get `6` and `25`, respectively.
-   *
-   * This method is also called for `==` equality.
-   */
-  [Symbol.toPrimitive](): any;
-  /** Realize the array and return it as data. */
-  data(): Promise<DataArray>;
-  /**
-   * Wait for this array to finish evaluation.
-   *
-   * Operations and data loading in jax-js are lazy, so this function ensures
-   * that pending operations are dispatched and fully executed before it
-   * returns.
-   *
-   * If you are mapping from `data()` or `dataSync()`, it will also trigger
-   * dispatch of operations as well.
-   *
-   * **Note:** `jax.blockUntilReady()` is a higher-level API, it calls this
-   * asynchronously for multiple arrays.
-   */
-  blockUntilReady(): Promise<Array>;
-  /**
-   * Realize the array and return it as data. This is a sync variant and not
-   * recommended for performance reasons, as it will block rendering.
-   */
-  dataSync(): DataArray;
-  /**
-   * Convert this array into a JavaScript object.
-   *
-   * This is a blocking operation that will compile all of the shaders and wait
-   * for execution to complete, synchronously. No other JavaScript code on the
-   * site will be run during shader execution.
-   *
-   * To avoid blocking, prefer `jsAsync()` when possible.
-   */
-  js(): any;
-  /** Convert this array into a JavaScript object, asynchronously. */
-  jsAsync(): Promise<any>;
-  /**
-   * Copy an element of an array to a numeric scalar and return it.
-   *
-   * Throws an error if the array does not have a single element. The array must
-   * either be rank-0, or all dimensions of the shape are 1.
-   */
-  item(): number;
-  /** @private Internal plumbing method for Array / Tracer ops. */
-  static _implRules(): typeof implRules;
-  /** @private */
-  _realizeSource(): number;
-  /** @private Put this array on a new backend, asynchronously. */
-  _put(backend: Backend): Promise<Array>;
-  /** @private Put this array on a new backend, synchronously. */
-  _putSync(backend: Backend): Array;
-}
-/** Constructor for creating a new array from data. */
-declare function array(values: Array | DataArray | RecursiveArray<number> | RecursiveArray<boolean>, {
-  shape,
-  dtype,
-  device
-}?: {
-  shape?: number[];
-} & DTypeAndDevice): Array;
-/** If x is a value, lift it into an array, otherwise leave it be. */
-
-type ImplRule<P extends Primitive> = (tracers: Array[], params: PrimitiveParams<P>) => Array[];
-declare const implRules: { [P in Primitive]: ImplRule<P> };
-/** Return a new array of given shape and type, filled with zeros. */
-declare function zeros(shape: number[], {
-  dtype,
-  device
-}?: DTypeAndDevice): Array;
-/** Return a new array of given shape and type, filled with ones. */
-declare function ones(shape: number[], {
-  dtype,
-  device
-}?: DTypeAndDevice): Array;
-/** Return a new array of given shape and type, filled with `fill_value`. */
-declare function full(shape: number[], fillValue: number | boolean | Array, {
-  dtype,
-  device
-}?: DTypeAndDevice): Array;
+declare const atan2: 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;
 /**
- * Create an identity matrix.
+ * Return the largest integer smaller or equal to the division of the inputs.
  *
- * If numCols is not provided, it defaults to numRows, i.e., a square identity
- * matrix with ones on the diagonal.
+ * 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 eye(numRows: number, numCols?: number, {
-  dtype,
-  device
-}?: DTypeAndDevice): Array;
-/** Return the identity matrix, with ones on the main diagonal. */
-declare function identity$1(n: number, {
-  dtype,
-  device
-}?: DTypeAndDevice): Array;
+declare function floorDivide(x: ArrayLike, y: ArrayLike): Array;
 /**
- * Return evenly spaced values within a given interval.
+ * @function
+ * Calculate element-wise floating-point modulo operation.
+ */
+declare const fmod: OwnedFunction<(x: ArrayLike, y: ArrayLike) => Array>;
+/**
+ * @function
+ * 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.
  *
- * This can be called with a varying number of arguments, just like the range()
- * builtin function in Python.
+ * Equivalent to `[floorDivide(x, y), remainder(x, y)]`.
  *
- * - `arange(stop)` is equivalent to `arange(0, stop, 1)`.
- * - `arange(start, stop)` is equivalent to `arange(start, stop, 1)`.
- * - `arange(start, stop, step)` creates an array starting at `start`, ending
- *   before `stop`, with a step size of `step`.
+ * @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;
+/**
+ * Compute `x1 * 2 ** x2` as a standard multiplication and exponentiation.
  *
- * Defaults to an integer data type. This can produce unintended results when
- * using a non-integer step, so prefer linspace() in those cases.
+ * This is the inverse of `frexp()`.
  */
-declare function arange(start: number, stop?: number, step?: number, {
-  dtype,
-  device
-}?: DTypeAndDevice): Array;
+declare function ldexp(x1: ArrayLike, x2: ArrayLike): Array;
 /**
- * Return an array with ones on and below the diagonal and zeros elsewhere.
+ * Decompose floating-point values into mantissa and two's exponent.
  *
- * 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.
+ * The mantissa is returned in the range `(-1, 1)` with magnitude `>= 0.5` if
+ * `x != 0`, and the exponent is an integer such that
+ * `x = mantissa * 2**exponent`.
  */
-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;
+declare function frexp(x: ArrayLike): [Array, Array];
+/** Calculate `2**p` for all p in the input array. */
+declare function exp2(p: ArrayLike): Array;
+/** Return the base-2 logarithm of x, element-wise. */
+declare function log2(x: ArrayLike): Array;
+/** Return the base-10 logarithm of x, element-wise. */
+declare function log10(x: ArrayLike): Array;
+/** Calculate `exp(x) - 1` element-wise. */
+declare function expm1(x: ArrayLike): Array;
+/** Calculate the natural logarithm of `1 + x` element-wise. */
+declare function log1p(x: ArrayLike): Array;
+/** Convert angles from degrees to radians. */
+declare function deg2rad(x: ArrayLike): Array;
+/** @function Alias of `jax.numpy.deg2rad()`. */
+declare const radians: typeof deg2rad;
+/** Convert angles from radians to degrees. */
+declare function rad2deg(x: ArrayLike): Array;
+/** @function Alias of `jax.numpy.rad2deg()`. */
+declare const degrees: typeof rad2deg;
 /**
- * Return evenly spaced numbers over a specified interval.
- *
- * Returns _num_ evenly spaced samples, calculated over the interval
- * [`start`, `stop`]. The endpoint `stop` is included in the result by default,
- * but this is controlled by the `endpoint` parameter.
- *
- * The default data type is Float32. Use arange() for integer steps.
+ * @function
+ * Computes first array raised to power of second array, element-wise.
  */
-declare function linspace(start: number, stop: number, num?: number, endpoint?: boolean, {
-  dtype,
-  device
-}?: DTypeAndDevice): Array;
+declare const power: OwnedFunction<(x1: ArrayLike, x2: ArrayLike) => Array>;
+/** @function Calculate the element-wise cube root of the input array. */
+declare const cbrt: OwnedFunction<(x: ArrayLike) => 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).
+ * @function
+ * Calculate element-wise hyperbolic sine of input.
  *
- * @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.
+ * `sinh(x) = (exp(x) - exp(-x)) / 2`
  */
-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 };
-}
+declare const sinh: OwnedFunction<(x: ArrayLike) => Array>;
 /**
- * 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]]
+ * @function
+ * Calculate element-wise hyperbolic cosine of input.
  *
- * // Upper Cholesky factorization:
- * const U = lax.linalg.cholesky(x, { upper: true });
- * // U ≈ [[1.4142135, 0.70710677], [0, 1.2247449]]
- * ```
+ * `cosh(x) = (exp(x) + exp(-x)) / 2`
  */
-declare function cholesky(a: ArrayLike, {
-  upper
-}?: {
-  upper?: boolean;
-}): Array;
+declare const cosh: OwnedFunction<(x: ArrayLike) => 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";
+ * @function
+ * Calculate element-wise hyperbolic tangent of input.
  *
- * 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]
- * ```
+ * `tanh(x) = sinh(x)/cosh(x) = (exp(x) - exp(-x)) / (exp(x) + exp(-x))`
  */
-declare function lu(x: ArrayLike): [Array, Array, Array];
+declare const tanh: OwnedFunction<(x: ArrayLike) => 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]);
+ * @function
+ * Calculate element-wise inverse hyperbolic sine of input.
  *
- * // Solve L @ x = b
- * const x = lax.linalg.triangularSolve(L, b, { leftSide: true, lower: true });
- * // x = [[2.], [5./3.]]
- * ```
+ * `arcsinh(x) = ln(x + sqrt(x^2 + 1))`
  */
-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, lax_linalg_d_exports as linalg, reduceWindow, stopGradient };
-}
+declare const arcsinh: OwnedFunction<(x: ArrayLike) => Array>;
 /**
- * Dimension numbers for general `dot()` primitive.
+ * @function
+ * Calculate element-wise inverse hyperbolic cosine of input.
  *
- * Contracting dimensions act as a tensor contraction (reduction) along the
- * given axis. They must be the same size in both operands. Batch dimensions
- * are treated as vectorized, leading batch dimensions.
+ * `arccosh(x) = ln(x + sqrt(x^2 - 1))`
+ */
+declare const arccosh: OwnedFunction<(x: ArrayLike) => Array>;
+/**
+ * @function
+ * Calculate element-wise inverse hyperbolic tangent of input.
  *
- * The return value has a shape where the first dimensions are shared batch
- * dimensions, followed by `lhs` non-contracting dimensions, followed by
- * `rhs` non-contracting dimensions.
+ * `arctanh(x) = 0.5 * ln((1 + x) / (1 - x))`
  */
-type DotDimensionNumbers = {
-  lhsContractingDims?: number[];
-  rhsContractingDims?: number[];
-  lhsBatchDims?: number[];
-  rhsBatchDims?: number[];
-};
+declare const arctanh: OwnedFunction<(x: ArrayLike) => Array>;
 /**
- * General dot product/contraction operator.
+ * Compute the variance of an array.
  *
- * Prefer higher-level functions like `jax.numpy.dot()`, `jax.numpy.matmul()`,
- * `jax.numpy.tensordot(), and `jax.numpy.einsum()` where possible.
+ * The variance is computed for the flattened array by default, otherwise over
+ * the specified axis.
+ *
+ * If `correction` is provided, the divisor in calculation is `N - correction`,
+ * where `N` represents the number of elements (e.g., for Bessel's correction).
  */
-declare function dot(lhs: Array, rhs: Array, {
-  lhsContractingDims: lc,
-  rhsContractingDims: rc,
-  lhsBatchDims: lb,
-  rhsBatchDims: rb
-}?: DotDimensionNumbers): Array;
-type PaddingType = "VALID" | "SAME" | "SAME_LOWER" | Pair[];
+declare function var_(x: ArrayLike, axis?: Axis, opts?: {
+  mean?: ArrayLike;
+  correction?: number;
+} & ReduceOpts): Array;
 /**
- * General n-dimensional convolution operator, with optional dilation.
+ * Compute the standard deviation of an array.
  *
- * The semantics of this operation mimic the `jax.lax.conv_general_dilated`
- * function in JAX, which wraps XLA's general convolution operator.
+ * The standard deviation is computed for the flattened array by default,
+ * otherwise over the specified axis.
  *
- * Grouped convolutions are not supported right now.
+ * If `correction` is provided, the divisor in calculation is `N - correction`,
+ * where `N` represents the number of elements (e.g., for Bessel's correction).
  */
-declare function convGeneralDilated(lhs: Array, rhs: Array, windowStrides: number[], padding: PaddingType, {
-  lhsDilation,
-  rhsDilation,
-  featureGroupCount
+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
 }?: {
-  lhsDilation?: number[];
-  rhsDilation?: number[];
-  featureGroupCount?: number;
+  rowvar?: boolean;
 }): Array;
-/** Convenience wrapper around `convGeneralDilated`. */
-declare function convWithGeneralPadding(lhs: Array, rhs: Array, windowStrides: number[], padding: PaddingType, lhsDilation?: number[], rhsDilation?: number[]): Array;
-/** Convenience wrapper around `convGeneralDilated`. */
-declare function conv(lhs: Array, rhs: Array, windowStrides: number[], padding: PaddingType): Array;
-/** Reduce a computation over padded windows. */
-declare function reduceWindow(operand: Array, computation: (x: Array) => Array, windowDimensions: number[], windowStrides?: number[]): Array;
-/** The error function: `erf(x) = 2/sqrt(pi) * int[0..x] exp(-t^2) dt`. */
-declare function erf(x: 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). */
+declare function isnan(x: ArrayLike): Array;
+/** Test element-wise for negative infinity, return bool array. */
+declare function isneginf(x: ArrayLike): Array;
+/** Test element-wise for positive infinity, return bool array. */
+declare function isposinf(x: ArrayLike): Array;
 /**
- * The complementary error function: `erfc(x) = 1 - erf(x)`.
+ * Replace NaN and infinite entries in an array.
  *
- * This function is more accurate than `1 - erf(x)` for large values of `x`,
- * where `erf(x)` is very close to 1.
+ * By default, NaNs are replaced with `0.0`, and infinities are are substituted
+ * with the corresponding maximum or minimum finite values.
  */
-declare function erfc(x: ArrayLike): Array;
+declare function nanToNum(x: ArrayLike, {
+  nan,
+  posinf,
+  neginf
+}?: {
+  nan?: ArrayLike;
+  posinf?: ArrayLike | null;
+  neginf?: ArrayLike | null;
+}): Array;
 /**
- * Stops gradient computation.
- *
- * Behaves as the identity function but prevents the flow of gradients during
- * forward or reverse-mode automatic differentiation.
+ * @function
+ * Test element-wise for finite values (not infinity or NaN).
  */
-declare function stopGradient(x: ArrayLike): Array;
+declare const isfinite: OwnedFunction<(x: ArrayLike) => Array>;
 declare namespace nn_d_exports {
-  export { celu, elu, gelu, glu, hardSigmoid, hardSilu, hardSilu as hardSwish, hardTanh, identity, leakyRelu, logSigmoid, logSoftmax, logmeanexp, logsumexp, mish, oneHot, relu, relu6, selu, sigmoid, silu, softSign, softmax, softplus, sparsePlus, sparseSigmoid, squareplus, standardize, silu as swish };
+  export { celu, dotProductAttention, elu, gelu, glu, hardSigmoid, hardSilu, hardSilu as hardSwish, hardTanh, identity, leakyRelu, logSigmoid, logSoftmax, logmeanexp, logsumexp, mish, oneHot, relu, relu6, selu, sigmoid, silu, softSign, softmax, softplus, sparsePlus, sparseSigmoid, squareplus, standardize, silu as swish };
 }
 /**
  * Rectified Linear Unit (ReLU) activation function:
@@ -2435,6 +2513,56 @@ declare function standardize(x: ArrayLike, axis?: Axis, opts?: {
  * ```
  */
 declare function oneHot(x: Array, numClasses: number): Array;
+/**
+ * Scaled dot product attention (SDPA).
+ *
+ * Computes `softmax((Q @ K^T) / sqrt(d) + bias) @ V`, where `Q` is the query,
+ * `K` is the key, `V` is the value, and `d` is the dimensionality of each key
+ * and query vector.
+ *
+ * Multi-query attention is applied when input `key` and `value` tensors have
+ * fewer heads than `query`.
+ *
+ * We use the following uppercase letters to denote array shapes:
+ * - `B` = batch size
+ * - `S` = length of key/value sequences (source)
+ * - `L` = length of query sequences
+ * - `N` = number of attention heads
+ * - `H` = dimensionality of each attention head
+ * - `K` = number of key/value heads (for grouped-query attention)
+ *
+ * The batch size `B` may be omitted, which is equivalent to `B = 1`. In this
+ * case it must be omitted from all inputs.
+ *
+ * @param query - Query array; shape `[B, L, N, H]`
+ * @param key - Key array; shape `[B, S, K, H]`
+ * @param value - Value array; same shape as `key`
+ * @param opts.bias - Optional bias to add to the attention logits; shape
+ *   `[B, N, L, S]` or broadcastable to it.
+ * @param opts.mask - Optional mask to apply to the attention logits; should be
+ *   a boolean array broadcastable to `[B, N, L, S]`, where `true` indicates
+ *   the element should take part in attention.
+ * @param opts.scale - Scaling factor override, default is `1 / sqrt(H)`.
+ * @param opts.isCausal - If true, applies a casual mask.
+ * @param opts.querySeqLengths - Optional sequence lengths for the queries;
+ *   shape `(B,)`. Taken from the beginning of the tensor.
+ * @param opts.keyValueSeqLengths - Optional sequence lengths for the keys and
+ *   values; shape `(B,)`. Taken from the beginning of the tensor.
+ * @param opts.localWindowSize - If specified, applies a local attention window
+ *   of the given size. Can be a single number or a tuple `[left, right]`.
+ *
+ * @returns The result of the attention operation; shape is the same as query
+ *   `[B, L, N, H]`, or `[L, N, H]` if `B` is omitted.
+ */
+declare function dotProductAttention(query: ArrayLike, key: ArrayLike, value: ArrayLike, opts?: {
+  bias?: ArrayLike;
+  mask?: ArrayLike;
+  scale?: number;
+  isCausal?: boolean;
+  querySeqLengths?: ArrayLike;
+  keyValueSeqLengths?: ArrayLike;
+  localWindowSize?: number | [number, number];
+}): Array;
 declare namespace random_d_exports {
   export { bernoulli, bits, cauchy, exponential, gumbel, key, laplace, multivariateNormal, normal, split, uniform };
 }
@@ -2526,7 +2654,9 @@ declare const logit: OwnedFunction<(x: ArrayLike) => Array>;
  * @function
  * Compute the forward-mode Jacobian-vector product for a function.
  */
-declare const jvp: <F extends (...args: any[]) => JsTree<Array>>(f: F, primals: MapJsTree<Parameters<F>, Array, ArrayLike>, tangents: MapJsTree<Parameters<F>, Array, ArrayLike>) => [ReturnType<F>, ReturnType<F>];
+declare const jvp: <F extends (...args: any[]) => JsTree<Array>, HA extends boolean = false>(f: F, primals: MapJsTree<Parameters<F>, Array, ArrayLike>, tangents: MapJsTree<Parameters<F>, Array, ArrayLike>, opts?: {
+  hasAux?: HA;
+}) => HA extends true ? ReturnType<F> extends [infer Out, infer Aux] ? [Out, Out, Aux] : never : [ReturnType<F>, ReturnType<F>];
 /**
  * @function
  * Vectorize an operation on a batched axis for one or more inputs.
@@ -2568,28 +2698,100 @@ declare const jit: <F extends (...args: any[]) => JsTree<Array>>(f: F, opts?: Ji
  * Produce a local linear approximation to a function at a point using jvp() and
  * partial evaluation.
  */
-declare const linearize: <F extends (...args: any[]) => JsTree<Array>>(f: F, ...primals: MapJsTree<Parameters<F>, Array, ArrayLike>) => [ReturnType<F>, (...tangents: MapJsTree<Parameters<F>, Array, ArrayLike>) => ReturnType<F>];
+declare const linearize: <F extends (...args: any[]) => JsTree<Array>, HA extends boolean = false>(f: F, primals: MapJsTree<Parameters<F>, Array, ArrayLike>, opts?: {
+  hasAux?: HA;
+}) => HA extends true ? ReturnType<F> extends [infer Out, infer Aux] ? [Out, OwnedFunction<(...tangents: MapJsTree<Parameters<F>, Array, ArrayLike>) => Out>, Aux] : never : [ReturnType<F>, OwnedFunction<(...tangents: MapJsTree<Parameters<F>, Array, ArrayLike>) => ReturnType<F>>];
 /**
  * @function
  * Calculate the reverse-mode vector-Jacobian product for a function.
+ *
+ * The return value is a tuple of `[out, vjpFn]`, where `out` is the output of
+ * `f(primals)`, and `vjpFn` is a function that takes in cotangents for each
+ * output and returns the cotangents for each input.
+ *
+ * When `{ hasAux: true }` is passed, the function `f` is expected to return an
+ * `[out, aux]` tuple, and `vjp` returns `[out, vjpFn, aux]`.
+ *
+ * @example
+ * ```ts
+ * const [y, vjpFn] = vjp(f, [x]);
+ *
+ * // With hasAux
+ * const [y, vjpFn, aux] = vjp(f, [x], { hasAux: true });
+ * ```
  */
-declare const vjp: <F extends (...args: any[]) => JsTree<Array>>(f: F, ...primals: MapJsTree<Parameters<F>, Array, ArrayLike>) => [ReturnType<F>, (cotangents: MapJsTree<ReturnType<F>, Array, ArrayLike>) => MapJsTree<Parameters<F>, ArrayLike, Array>];
+declare const vjp: <F extends (...args: any[]) => JsTree<Array>, const HA extends boolean = false>(f: F, primals: MapJsTree<Parameters<F>, Array, ArrayLike>, opts?: {
+  hasAux?: HA;
+}) => HA extends true ? ReturnType<F> extends [infer Out, infer Aux] ? [Out, OwnedFunction<(cotangents: MapJsTree<Out, Array, ArrayLike>) => MapJsTree<Parameters<F>, ArrayLike, Array>>, Aux] : never : [ReturnType<F>, OwnedFunction<(cotangents: MapJsTree<ReturnType<F>, Array, ArrayLike>) => MapJsTree<Parameters<F>, ArrayLike, Array>>];
+/** @inline */
+type GradOutputType<I, F extends (...args: any[]) => any> = MapJsTree<I extends undefined ? Parameters<F>[0] : I extends number ? Parameters<F>[I] : I extends number[] ? { [K in keyof I]: I[K] extends number ? Parameters<F>[I[K]] : never } : never, ArrayLike, Array>;
 /**
  * @function
  * Compute the gradient of a scalar-valued function `f` with respect to its
  * first argument.
+ *
+ * Pass in different `argnums` to differentiate with respect to other
+ * arguments. If a tuple is provided, the return value will be a tuple of
+ * gradients corresponding to each argument index.
+ *
+ * When `{ hasAux: true }` is passed, the function `f` is expected to return a
+ * `[out, aux]` tuple, and the return value will be `[gradient, aux]`.
+ *
+ * @example
+ * ```ts
+ * const gradient = grad(f)(x);
+ *
+ * // With `argnums`
+ * const [gradientX, gradientZ] = grad(f, { argnums: [0, 2] })(x, y, z);
+ *
+ * // With `hasAux`
+ * const [gradient, aux] = grad(f, { hasAux: true })(x);
+ * ```
  */
-declare const grad: <F extends (...args: any[]) => JsTree<Array>>(f: F) => (...primals: MapJsTree<Parameters<F>, Array, ArrayLike>) => MapJsTree<Parameters<F>[0], ArrayLike, Array>;
+declare const grad: <F extends (...args: any[]) => JsTree<Array>, const I extends undefined | number | number[] = undefined, const HA extends boolean = false>(f: F, opts?: Omit<GradOpts, "argnums" | "hasAux"> & {
+  argnums?: I;
+  hasAux?: HA;
+}) => (...primals: MapJsTree<Parameters<F>, Array, ArrayLike>) => HA extends true ? ReturnType<F> extends [any, infer Aux] ? [GradOutputType<I, F>, Aux] : never : GradOutputType<I, F>;
 /**
  * @function
  * Create a function that evaluates both `f` and the gradient of `f`.
+ *
+ * When `{ hasAux: true }` is passed, the function `f` is expected to return an
+ * `[out, aux]` tuple, and the return value will be `[[out, aux], gradient]`.
+ *
+ * @example
+ * ```ts
+ * // Without hasAux
+ * const [value, gradient] = valueAndGrad(f)(x);
+ *
+ * // With hasAux
+ * const [[value, aux], gradient] = valueAndGrad(f, { hasAux: true })(x);
+ * ```
  */
-declare const valueAndGrad: <F extends (...args: any[]) => JsTree<Array>>(f: F) => (...primals: MapJsTree<Parameters<F>, Array, ArrayLike>) => [ReturnType<F>, MapJsTree<Parameters<F>[0], ArrayLike, Array>];
+declare const valueAndGrad: <F extends (...args: any[]) => JsTree<Array>, const I extends undefined | number | number[] = undefined, const HA extends boolean = false>(f: F, opts?: Omit<GradOpts, "argnums"> & {
+  argnums?: I;
+  hasAux?: HA;
+}) => (...primals: MapJsTree<Parameters<F>, Array, ArrayLike>) => [ReturnType<F>, GradOutputType<I, F>];
 /**
  * @function
  * Compute the Jacobian evaluated row-by-row by reverse-mode AD.
  */
 declare const jacrev: typeof jacfwd;
+/**
+ * @function
+ * Compute the Hessian matrix of a scalar-valued function.
+ *
+ * The Hessian is the matrix of second-order partial derivatives of a function.
+ * This is implemented as `jacfwd(grad(f))`.
+ *
+ * @example
+ * ```ts
+ * const f = (x: np.Array) => np.sum(x.ref.mul(x.ref).mul(x)); // x^3
+ * const H = hessian(f)(np.array([1, 2, 3]));
+ * // H[i,j] = d^2f / dx_i dx_j
+ * ```
+ */
+declare const hessian: <F extends (x: Array) => Array>(f: F) => (...args: MapJsTree<Parameters<F>, Array, ArrayLike>) => ReturnType<F>;
 /**
  * Wait until all `Array` leaves are ready by calling `Array.blockUntilReady()`.
  *
@@ -2612,4 +2814,4 @@ declare function blockUntilReady<T extends JsTree<any>>(x: T): Promise<T>;
  */
 declare function devicePut<T extends JsTree<any>>(x: T, device?: Device): Promise<MapJsTree<T, number | boolean, Array>>;
 //#endregion
-export { Array, ClosedJaxpr, DType, type Device, Jaxpr, type JsTree, type JsTreeDef, type OwnedFunction, blockUntilReady, defaultDevice, devicePut, devices, grad, init, jacfwd, jacrev as jacobian, jacrev, jit, jvp, lax_d_exports as lax, linearize, makeJaxpr, nn_d_exports as nn, numpy_d_exports as numpy, random_d_exports as random, scipy_special_d_exports as scipySpecial, setDebug, tree_d_exports as tree, valueAndGrad, vjp, vmap };
+export { Array, ClosedJaxpr, DType, type Device, Jaxpr, type JsTree, type JsTreeDef, type OwnedFunction, blockUntilReady, defaultDevice, devicePut, devices, grad, hessian, init, jacfwd, jacrev as jacobian, jacrev, jit, jvp, lax_d_exports as lax, linearize, makeJaxpr, nn_d_exports as nn, numpy_d_exports as numpy, random_d_exports as random, scipy_special_d_exports as scipySpecial, setDebug, tree_d_exports as tree, valueAndGrad, vjp, vmap };