@jax-js/jax 0.1.6 → 0.1.7

package/dist/index.d.ts

@@ -538,1757 +538,1780 @@ 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, swapaxes, 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;
-/**
- * @function
- * Add padding (zeros) to an array.
- *
- * 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.
- */
-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;
-/**
- * Test whether all array elements along a given axis evaluate to True.
- *
- * Returns a boolean array with the same shape as `a` with the specified axis
- * removed. If axis is None, returns a scalar.
- */
-declare function all(a: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
-/**
- * Test whether any array element along a given axis evaluates to True.
- *
- * Returns a boolean array with the same shape as `a` with the specified axis
- * removed. If axis is None, returns a scalar.
- */
-declare function any(a: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
-/** Return the peak-to-peak range along a given axis (`max - min`). */
-declare function ptp(a: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
-/** Compute the average of the array elements along the specified axis. */
-declare function mean(a: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
-/**
- * Returns the indices of the minimum values along an axis.
- *
- * By default, index is into the flatted array, otherwise it is along the
- * specified axis.
- */
-declare function argmin(a: ArrayLike, axis?: number, opts?: ReduceOpts): Array;
-/**
- * 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 function argmax(a: ArrayLike, axis?: number, opts?: ReduceOpts): Array;
-/**
- * Cumulative sum of elements along an axis.
- *
- * Currently this function is `O(n^2)`, we'll improve this later on with a
- * two-phase parallel reduction algorithm.
- */
-declare function cumsum(a: ArrayLike, axis?: number): Array;
-/** Reverse the elements in an array along the given axes. */
-declare function flip(x: ArrayLike, axis?: Axis): Array;
-/**
- * Split an array into multiple sub-arrays along an axis.
- *
- * @param a - The input array to split.
- * @param indicesOrSections - If an integer, it indicates the number of equal
- * sections to create along the specified axis. If a list of integers, it
- * specifies the indices at which to split the array.
- * @param axis - The axis along which to split the array. Default is 0.
- */
-declare function split$1(a: ArrayLike, indicesOrSections: number | number[], axis?: number): Array[];
-/**
- * Join a sequence of arrays along an existing axis.
- *
- * The arrays must have the same shape, except in the dimension corresponding to
- * `axis` (the first, by default).
- *
- * No scalars can be passed to this function, as the axis is then ambiguous.
- */
-declare function concatenate(xs: Array[], axis?: number): Array;
-/**
- * Join a sequence of arrays along a new axis.
- *
- * 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.
- *
- * All shapes must have the same shape.
- */
-declare function stack(xs: ArrayLike[], axis?: number): 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).
- */
-declare function hstack(xs: ArrayLike[]): Array;
-/**
- * Vertically stack arrays. Inputs are promoted to rank at least 2, then
- * concatenated along axis 0.
- */
-declare function vstack(xs: ArrayLike[]): Array;
-/**
- * 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;
+
+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. */
+
 /**
- * Expand the shape of an array by inserting new axes of length 1.
- *
- * @param a - Input array.
- * @param axis - Position(s) in the expanded axes where the new axis (or axes)
- *   is placed. Can be a single integer or an array of integers.
- * @returns Array with the number of dimensions increased.
+ * An executable operation that will be dispatched to the backend.
  *
- * @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]
- * ```
+ * This holds a reference to all input buffers used in the operation. After the
+ * operation is dispatched, the references should be released.
  */
-declare function expandDims(a: ArrayLike, axis: number | number[]): 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>;
+};
 /**
- * Repeat each element of an array after themselves.
+ * A multidimensional numeric array with data stored on CPU or GPU.
  *
- * 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;
-/**
- * Construct an array by repeating A the number of times given by reps.
+ * This is the library's core data type. Equivalent to `jax.Array` from JAX, or
+ * `torch.Tensor`.
  *
- * 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.
+ * 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 tile(a: ArrayLike, reps: number | number[]): 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;
 /**
- * Broadcast an array to a shape, with NumPy-style broadcasing rules.
+ * Create an identity matrix.
  *
- * In other words, this lets you append axes to the left, and/or expand
- * dimensions where the shape is 1.
+ * If numCols is not provided, it defaults to numRows, i.e., a square identity
+ * matrix with ones on the diagonal.
  */
-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[];
+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;
 /**
- * Return specified diagonals.
+ * Return evenly spaced values within a given interval.
  *
- * 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 can be called with a varying number of arguments, just like the range()
+ * builtin function in Python.
  *
- * 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 function diagonal(a: ArrayLike, offset?: number, axis1?: number, axis2?: number): Array;
-/**
- * Extract a diagonal or construct a diagonal array.
+ * - `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`.
  *
- * 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.
+ * 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 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 arange(start: number, stop?: number, step?: number, {
+  dtype,
+  device
+}?: DTypeAndDevice): Array;
 /**
- * Return a sorted copy of an array.
+ * Return an array with ones on and below the diagonal and zeros elsewhere.
  *
- * 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.
+ * 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 sort(a: ArrayLike, axis?: number): 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;
 /**
- * Return indices that would sort an array. This may be an unstable sorting
- * algorithm; it need not preserve order of indices in ties.
+ * Return evenly spaced numbers over a specified interval.
  *
- * Returns an array of `int32` indices.
+ * 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 array is sorted along a specified axis (the last by default).
+ * The default data type is Float32. Use arange() for integer steps.
  */
-declare function argsort(a: ArrayLike, axis?: number): Array;
+declare function linspace(start: number, stop: number, num?: number, endpoint?: boolean, {
+  dtype,
+  device
+}?: DTypeAndDevice): Array;
 /**
- * Take elements from an array along an axis.
+ * Return numbers spaced evenly on a log scale.
  *
- * This is equivalent to advanced indexing with integer indices over that
- * numbered axis. By default, the flattened array is used.
- */
-declare function take(a: ArrayLike, indices: ArrayLike, axis?: number | null): Array;
-/** Return if two arrays are element-wise equal within a tolerance. */
-declare function allclose(actual: Parameters<typeof array>[0], expected: Parameters<typeof array>[0], options?: {
-  rtol?: number;
-  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;
-/**
- * Compute the tensor dot product of two N-dimensional arrays.
+ * In linear space, the sequence starts at `base ** start` and ends at
+ * `base ** stop` (see `endpoint` below).
  *
- * 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.
+ * @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 tensordot(x: ArrayLike, y: ArrayLike, axes?: number | [number[], 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 };
+}
 /**
- * Einstein summation with string subscripts.
+ * 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 { numpy as np } from "@jax-js/jax";
+ * import { lax, 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]
+ * 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 einsum(subscripts: string, ...operands: ArrayLike[]): Array;
+declare function cholesky$1(a: ArrayLike, {
+  upper
+}?: {
+  upper?: boolean;
+}): Array;
 /**
- * Einstein summation alternating between arrays and numeric indices.
+ * 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 { numpy as np } from "@jax-js/jax";
+ * import { lax, numpy as np } from "@jax-js/jax";
  *
- * const a = np.ones([2, 3]);
- * const b = np.ones([3]);
- * np.einsum(a, [0, 1], b, [1]); // Shape [2]
+ * 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 einsum(...args: (ArrayLike | number[])[]): Array;
+declare function lu(x: ArrayLike): [Array, Array, Array];
 /**
- * Compute the inner product of two arrays.
+ * Solve a triangular linear system.
  *
- * Unlike `jax.numpy.matmul()` or `jax.numpy.dot()`, this always performs a
- * contraction on the last axis.
+ * Solves `a @ x = b` (if leftSide=true) or `x @ a = b` (if leftSide=false)
+ * where `a` is a triangular matrix.
  *
- * Returned array has shape `[...x.shape[:-1], ...y.shape[:-1]]`.
- */
-declare function inner(x: ArrayLike, y: ArrayLike): Array;
-/**
- * Compute the outer product of two arrays.
+ * @example
+ * ```ts
+ * import { lax, numpy as np } from "@jax-js/jax";
  *
- * If the input arrays are not 1D, they will be flattened. Returned array will
- * be of shape `[x.size, y.size]`.
+ * 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 outer(x: ArrayLike, y: ArrayLike): Array;
-/** Vector dot product of two arrays along a given axis. */
-declare function vecdot(x: ArrayLike, y: ArrayLike, {
-  axis
+declare function triangularSolve(a: ArrayLike, b: ArrayLike, {
+  leftSide,
+  lower,
+  transposeA,
+  unitDiagonal
 }?: {
-  axis?: number;
+  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 };
+}
 /**
- * Return the dot product of two vectors.
+ * Dimension numbers for general `dot()` primitive.
  *
- * Like vecdot() but flattens the arguments first into vectors.
+ * 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 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;
+type DotDimensionNumbers = {
+  lhsContractingDims?: number[];
+  rhsContractingDims?: number[];
+  lhsBatchDims?: number[];
+  rhsBatchDims?: number[];
+};
+/**
+ * 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 dot$1(lhs: Array, rhs: Array, {
+  lhsContractingDims: lc,
+  rhsContractingDims: rc,
+  lhsBatchDims: lb,
+  rhsBatchDims: rb
+}?: DotDimensionNumbers): Array;
+type PaddingType = "VALID" | "SAME" | "SAME_LOWER" | Pair[];
 /**
- * Return a tuple of coordinate matrices from coordinate vectors.
+ * General n-dimensional convolution operator, with optional dilation.
  *
- * 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.
+ * 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 meshgrid(xs: Array[], {
-  indexing
+declare function convGeneralDilated(lhs: Array, rhs: Array, windowStrides: number[], padding: PaddingType, {
+  lhsDilation,
+  rhsDilation,
+  featureGroupCount
 }?: {
-  indexing?: "xy" | "ij";
-}): Array[];
+  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;
 /**
- * Clip (limit) the values in an array.
+ * Convenience wrapper for calculating the N-d convolution "transpose".
  *
- * 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.
+ * 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:
  *
- * If either bound is undefined, it is ignored.
- */
-declare function clip(a: ArrayLike, min?: ArrayLike, max?: ArrayLike): Array;
-/**
- * Calculate the absolute value element-wise.
+ * - 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`.
  *
- * This is the same function as `jax.numpy.abs()`.
+ * 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 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 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;
 /**
- * Return the Hamming window of size M, a taper with a weighted cosine bell.
+ * The complementary error function: `erfc(x) = 1 - erf(x)`.
  *
- * `w(n) = 0.54 - 0.46 * cos(2πn/(M-1))` for `0 <= n <= M-1`.
+ * This function is more accurate than `1 - erf(x)` for large values of `x`,
+ * where `erf(x)` is very close to 1.
  */
-declare function hamming(M: number): Array;
+declare function erfc(x: ArrayLike): Array;
 /**
- * Return the Hann window of size M, a taper with a weighted cosine bell.
+ * Stops gradient computation.
  *
- * `w(n) = 0.5 - 0.5 * cos(2πn/(M-1))` for `0 <= n <= M-1`.
+ * Behaves as the identity function but prevents the flow of gradients during
+ * forward or reverse-mode automatic differentiation.
  */
-declare function hann(M: number): Array;
+declare function stopGradient(x: ArrayLike): Array;
+declare namespace numpy_fft_d_exports {
+  export { ComplexPair, fft, ifft };
+}
 /**
- * @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`.
+ * A pair of arrays representing real and imaginary part `a + bj`. Both arrays
+ * must have the same shape.
  */
-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;
+type ComplexPair = {
+  real: Array;
+  imag: Array;
+};
 /**
- * @function
- * Return the normalized sinc function.
+ * Compute a one-dimensional discrete Fourier transform.
  *
- * 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.
+ * 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.
  *
- * **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, the size of the axis must be a power of two.
  */
-declare const sinc: OwnedFunction<(x: ArrayLike) => Array>;
-/** Element-wise inverse cosine function (inverse of cos). */
-declare function acos(x: ArrayLike): 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 };
+}
 /**
- * @function
- * Return element-wise hypotenuse for the given legs of a right triangle.
+ * Compute the Cholesky decomposition of a (batched) positive-definite matrix.
  *
- * 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.
+ * This is like `jax.lax.linalg.cholesky()`, except with an option to symmetrize
+ * the input matrix, which is on by default.
  */
-declare const hypot: OwnedFunction<(x1: ArrayLike, x2: ArrayLike) => 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;
 /**
- * @function
- * Element-wise arc tangent of y/x with correct quadrant.
+ * Return the least-squares solution to a linear equation.
  *
- * Returns the angle in radians between the positive x-axis and the point (x, y).
- * The result is in the range [-π, π].
+ * For overdetermined systems, this finds the `x` that minimizes `norm(ax - b)`.
+ * For underdetermined systems, this finds the minimum-norm solution for `x`.
  *
- * 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)
+ * This currently uses Cholesky decomposition to solve the normal equations,
+ * under the hood. The method is not as robust as QR or SVD.
  *
- * The output is ill-defined when both x and y are zero.
+ * @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 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 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];
 /**
- * Return the largest integer smaller or equal to the division of the inputs.
- *
- * The result is always rounded towards negative infinity.
+ * Solve a linear system of equations.
  *
- * 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).
+ * 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 x - Dividend array.
- * @param y - Divisor array.
- * @returns Element-wise floor division of x by y.
+ * @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 floorDivide(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;
 /**
  * @function
- * Calculate element-wise floating-point modulo operation.
+ * Permute the dimensions of an array. Defaults to reversing the axis order.
  */
-declare const fmod: OwnedFunction<(x: ArrayLike, y: ArrayLike) => Array>;
+declare const transpose: (x: ArrayLike, perm?: number[]) => 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.
- *
- * Equivalent to `[floorDivide(x, y), remainder(x, y)]`.
- *
- * @param x - Dividend array.
- * @param y - Divisor array.
- * @returns Tuple of [quotient, remainder].
- */
-declare function divmod(x: ArrayLike, y: ArrayLike): [Array, Array];
-/** Round input to the nearest integer towards zero. */
-declare function trunc(x: ArrayLike): Array;
-/**
- * Compute `x1 * 2 ** x2` as a standard multiplication and exponentiation.
- *
- * This is the inverse of `frexp()`.
- */
-declare function ldexp(x1: ArrayLike, x2: ArrayLike): Array;
-/**
- * Decompose floating-point values into mantissa and two's exponent.
+ * Give a new shape to an array without changing its data.
  *
- * 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 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;
-/**
- * @function
- * Computes first array raised to power of second array, element-wise.
+ * One shape dimension can be -1. In this case, the value is inferred from the
+ * length of the array and remaining dimensions.
  */
-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 const reshape: (x: ArrayLike, shape: number[]) => Array;
 /**
  * @function
- * Calculate element-wise hyperbolic sine of input.
- *
- * `sinh(x) = (exp(x) - exp(-x)) / 2`
+ * Move axes of an array to new positions. Other axes retain original order.
  */
-declare const sinh: OwnedFunction<(x: ArrayLike) => Array>;
+declare const moveaxis: (x: ArrayLike, src: number, dst: number) => Array;
 /**
  * @function
- * Calculate element-wise hyperbolic cosine of input.
+ * Add padding (zeros) to an array.
  *
- * `cosh(x) = (exp(x) + exp(-x)) / 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 const cosh: OwnedFunction<(x: ArrayLike) => Array>;
+declare const pad: (x: ArrayLike, width: number | Pair | Pair[] | Record<number, Pair>) => Array;
 /**
  * @function
- * Calculate element-wise hyperbolic tangent of input.
- *
- * `tanh(x) = sinh(x)/cosh(x) = (exp(x) - exp(-x)) / (exp(x) + exp(-x))`
+ * Return the number of dimensions of an array. Does not consume array reference.
  */
-declare const tanh: OwnedFunction<(x: 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[];
 /**
  * @function
- * Calculate element-wise inverse hyperbolic sine of input.
- *
- * `arcsinh(x) = ln(x + sqrt(x^2 + 1))`
+ * Return an array of zeros with the same shape and type as a given array.
  */
-declare const arcsinh: OwnedFunction<(x: ArrayLike) => Array>;
+declare const zerosLike: (a: ArrayLike, dtype?: DType) => Array;
 /**
  * @function
- * Calculate element-wise inverse hyperbolic cosine of input.
- *
- * `arccosh(x) = ln(x + sqrt(x^2 - 1))`
+ * Return an array of ones with the same shape and type as a given array.
  */
-declare const arccosh: OwnedFunction<(x: ArrayLike) => Array>;
+declare const onesLike: (a: ArrayLike, dtype?: DType) => Array;
 /**
  * @function
- * Calculate element-wise inverse hyperbolic tangent of input.
- *
- * `arctanh(x) = 0.5 * ln((1 + x) / (1 - x))`
- */
-declare const arctanh: OwnedFunction<(x: ArrayLike) => Array>;
-/**
- * Compute the variance of an array.
- *
- * 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 var_(x: ArrayLike, axis?: Axis, opts?: {
-  mean?: ArrayLike;
-  correction?: number;
-} & ReduceOpts): Array;
-/**
- * Compute the standard deviation of an array.
- *
- * The standard deviation 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).
+ * Return a full array with the same shape and type as a given array.
  */
-declare function std(x: ArrayLike, axis?: Axis, opts?: {
-  mean?: ArrayLike;
-  correction?: number;
-} & ReduceOpts): Array;
-/** Estimate the sample covariance of a set of variables. */
-declare function cov(x: ArrayLike, y?: ArrayLike | null, {
-  rowvar
-}?: {
-  rowvar?: boolean;
-}): Array;
-/** Compute the Pearson correlation coefficients (in range `[-1, 1]`). */
-declare function corrcoef(x: ArrayLike, y?: ArrayLike): Array;
-/** Test element-wise for positive or negative infinity, return bool array. */
-declare function isinf(x: ArrayLike): Array;
-/** Test element-wise for NaN (Not a Number). */
-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;
+declare const fullLike: (a: ArrayLike, fillValue: number | boolean | Array, dtype?: DType) => Array;
 /**
- * @function
- * Test element-wise for finite values (not infinity or NaN).
+ * Return the number of elements in an array, optionally along an axis.
+ * Does not consume array reference.
  */
-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 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;
 /**
- * Check that the shapes and parameters passed to convolution are valid.
- * Expected shapes of the lhs and rhs of the convolution are:
- *
- * - `lhsShape = [*vmapDims, batchSize, inChannels, spatialDims...]`
- * - `rhsShape = [*vmapDims, outChannels, inChannels, kernelSize...]`
+ * Test whether any array element along a given axis evaluates to True.
  *
- * If the check succeeds, returns the output shape.
+ * Returns a boolean array with the same shape as `a` with the specified axis
+ * removed. If axis is None, returns a scalar.
  */
-//#endregion
-//#region src/frontend/jaxpr.d.ts
+declare function any(a: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
 /**
- * Function callback with an associated dispose() method.
+ * Test whether all array elements along a given axis evaluate to True.
  *
- * The dispose() method should be called to clean up any tracer resources needed
- * by the function after the last time it is called.
+ * Returns a boolean array with the same shape as `a` with the specified axis
+ * removed. If axis is None, returns a scalar.
  */
-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 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;
 /**
- * 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.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;
-};
+ * Returns the indices of the minimum values along an axis.
+ *
+ * By default, index is into the flatted array, otherwise it is along the
+ * specified axis.
+ */
+declare function argmin(a: ArrayLike, axis?: number, opts?: ReduceOpts): Array;
 /**
- * Push an interpreter onto the trace stack. Use this like:
- * `using main = newMain(...);`
+ * 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.
  */
-
-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 argmax(a: ArrayLike, axis?: number, opts?: ReduceOpts): Array;
 /**
- * Broadcast shapes and promote types with casting for two avals.
+ * Cumulative sum of elements along an axis.
  *
- * This implements the weak type behavior described in `promoteTypes()`, but not
- * implemented in that function as `weakType` is not passed.
+ * Currently this function is `O(n^2)`, we'll improve this later on with a
+ * two-phase parallel reduction algorithm.
  */
-
-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 function cumsum(a: ArrayLike, axis?: number): Array;
+/** Reverse the elements in an array along the given axes. */
+declare function flip(x: ArrayLike, axis?: Axis): Array;
+/**
+ * Split an array into multiple sub-arrays along an axis.
+ *
+ * @param a - The input array to split.
+ * @param indicesOrSections - If an integer, it indicates the number of equal
+ * sections to create along the specified axis. If a list of integers, it
+ * specifies the indices at which to split the array.
+ * @param axis - The axis along which to split the array. Default is 0.
+ */
+declare function split$1(a: ArrayLike, indicesOrSections: number | number[], axis?: number): Array[];
+/**
+ * Join a sequence of arrays along an existing axis.
+ *
+ * The arrays must have the same shape, except in the dimension corresponding to
+ * `axis` (the first, by default).
+ *
+ * No scalars can be passed to this function, as the axis is then ambiguous.
+ */
+declare function concatenate(xs: Array[], axis?: number): Array;
 /**
- * An executable operation that will be dispatched to the backend.
+ * Join a sequence of arrays along a new axis.
  *
- * This holds a reference to all input buffers used in the operation. After the
- * operation is dispatched, the references should be released.
+ * 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.
+ *
+ * All shapes must have the same shape.
  */
-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 function stack(xs: ArrayLike[], axis?: number): Array;
 /**
- * A multidimensional numeric array with data stored on CPU or GPU.
+ * 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 function hstack(xs: ArrayLike[]): Array;
+/**
+ * Vertically stack arrays. Inputs are promoted to rank at least 2, then
+ * concatenated along axis 0.
+ */
+declare function vstack(xs: ArrayLike[]): Array;
+/**
+ * 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.
  *
- * This is the library's core data type. Equivalent to `jax.Array` from JAX, or
- * `torch.Tensor`.
+ * @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.
  *
- * 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.
+ * @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 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 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;
+/**
+ * Construct an array by repeating A the number of times given by reps.
+ *
+ * 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 tile(a: ArrayLike, reps: number | number[]): Array;
+/**
+ * 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.
+ */
+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[];
+/**
+ * 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 function diagonal(a: ArrayLike, offset?: number, axis1?: number, axis2?: number): Array;
+/**
+ * Extract a diagonal or construct a diagonal array.
+ *
+ * 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 function diag(v: ArrayLike, k?: number): Array;
+/** Calculate the sum of the diagonal of an array along the given axes. */
+declare function trace(a: ArrayLike, offset?: number, axis1?: number, axis2?: number): Array;
+/**
+ * Return a sorted copy of an array.
+ *
+ * The array is sorted along a specified axis (the last by default). This may be
+ * an unstable sort, and it dispatches to device-specific implementation.
+ */
+declare function sort(a: ArrayLike, axis?: number): Array;
+/**
+ * Return indices that would sort an array. This may be an unstable sorting
+ * algorithm; it need not preserve order of indices in ties.
+ *
+ * Returns an array of `int32` indices.
+ *
+ * The array is sorted along a specified axis (the last by default).
+ */
+declare function argsort(a: ArrayLike, axis?: number): Array;
 /**
- * Create an identity matrix.
+ * Take elements from an array along an axis.
  *
- * If numCols is not provided, it defaults to numRows, i.e., a square identity
- * matrix with ones on the diagonal.
+ * This is equivalent to advanced indexing with integer indices over that
+ * numbered axis. By default, the flattened array is used.
  */
-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 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;
 /**
- * Return evenly spaced values within a given interval.
+ * Compute the tensor dot product of two N-dimensional arrays.
  *
- * This can be called with a varying number of arguments, just like the range()
- * builtin function in Python.
+ * 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 function tensordot(x: ArrayLike, y: ArrayLike, axes?: number | [number[], number[]]): Array;
+/**
+ * Einstein summation with string subscripts.
  *
- * - `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`.
+ * @example
+ * ```ts
+ * import { numpy as np } from "@jax-js/jax";
  *
- * Defaults to an integer data type. This can produce unintended results when
- * using a non-integer step, so prefer linspace() in those cases.
+ * const a = np.ones([2, 3]);
+ * const b = np.ones([3]);
+ * np.einsum("ij,j", a, b); // Shape [2]
+ * ```
  */
-declare function arange(start: number, stop?: number, step?: number, {
-  dtype,
-  device
-}?: DTypeAndDevice): Array;
+declare function einsum(subscripts: string, ...operands: ArrayLike[]): Array;
 /**
- * Return an array with ones on and below the diagonal and zeros elsewhere.
+ * Einstein summation alternating between arrays and numeric indices.
  *
- * 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.
+ * @example
+ * ```ts
+ * import { numpy as np } from "@jax-js/jax";
+ *
+ * const a = np.ones([2, 3]);
+ * const b = np.ones([3]);
+ * np.einsum(a, [0, 1], b, [1]); // Shape [2]
+ * ```
  */
-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 einsum(...args: (ArrayLike | number[])[]): Array;
 /**
- * Return evenly spaced numbers over a specified interval.
+ * Compute the inner product of two arrays.
  *
- * 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.
+ * Unlike `jax.numpy.matmul()` or `jax.numpy.dot()`, this always performs a
+ * contraction on the last axis.
  *
- * The default data type is Float32. Use arange() for integer steps.
+ * Returned array has shape `[...x.shape[:-1], ...y.shape[:-1]]`.
  */
-declare function linspace(start: number, stop: number, num?: number, endpoint?: boolean, {
-  dtype,
-  device
-}?: DTypeAndDevice): Array;
+declare function inner(x: ArrayLike, y: ArrayLike): Array;
 /**
- * Return numbers spaced evenly on a log scale.
+ * Compute the outer product of two arrays.
  *
- * In linear space, the sequence starts at `base ** start` and ends at
- * `base ** stop` (see `endpoint` below).
+ * 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
+}?: {
+  axis?: number;
+}): Array;
+/**
+ * Return the dot product of two vectors.
  *
- * @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.
+ * Like vecdot() but flattens the arguments first into vectors.
  */
-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, lu, triangularSolve };
-}
+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;
 /**
- * Compute the Cholesky decomposition of a symmetric positive-definite matrix.
+ * Return a tuple of coordinate matrices from coordinate vectors.
  *
- * The Cholesky decomposition of a matrix `A` is:
+ * 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[];
+/**
+ * Clip (limit) the values in an array.
  *
- * - A = L @ L^T  (for upper=false, default)
- * - A = U^T @ U  (for upper=true)
+ * 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.
  *
- * where `L` is a lower-triangular matrix and `U` is an upper-triangular matrix.
- * The input matrix must be symmetric and positive-definite.
+ * If either bound is undefined, it is ignored.
+ */
+declare function clip(a: ArrayLike, min?: ArrayLike, max?: ArrayLike): Array;
+/**
+ * Calculate the absolute value element-wise.
  *
- * @example
- * ```ts
- * import { lax, numpy as np } from "@jax-js/jax";
+ * This is the same function as `jax.numpy.abs()`.
+ */
+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;
+/**
+ * Return the Hamming window of size M, a taper with a weighted cosine bell.
  *
- * const x = np.array([[2., 1.], [1., 2.]]);
+ * `w(n) = 0.54 - 0.46 * cos(2πn/(M-1))` for `0 <= n <= M-1`.
+ */
+declare function hamming(M: number): Array;
+/**
+ * Return the Hann window of size M, a taper with a weighted cosine bell.
  *
- * // Lower Cholesky factorization (default):
- * const L = lax.linalg.cholesky(x);
- * // L ≈ [[1.4142135, 0], [0.70710677, 1.2247449]]
+ * `w(n) = 0.5 - 0.5 * cos(2πn/(M-1))` for `0 <= n <= M-1`.
+ */
+declare function hann(M: number): 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`.
+ */
+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.
  *
- * // Upper Cholesky factorization:
- * const U = lax.linalg.cholesky(x, { upper: true });
- * // U ≈ [[1.4142135, 0.70710677], [0, 1.2247449]]
- * ```
+ * 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 function cholesky(a: ArrayLike, {
-  upper
-}?: {
-  upper?: boolean;
-}): Array;
+declare const sinc: OwnedFunction<(x: ArrayLike) => Array>;
+/** Element-wise inverse cosine function (inverse of cos). */
+declare function acos(x: ArrayLike): Array;
 /**
- * LU decomposition with partial pivoting.
+ * @function
+ * Return element-wise hypotenuse for the given legs of a right triangle.
+ *
+ * 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 const hypot: OwnedFunction<(x1: ArrayLike, x2: ArrayLike) => 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 [-π, π].
+ *
+ * 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 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;
+/**
+ * Return the largest integer smaller or equal to the division of the inputs.
  *
- * 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.
+ * The result is always rounded towards negative infinity.
  *
- * @param x - A batch of matrices with shape `[..., m, n]`.
+ * 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).
  *
- * @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]`.
+ * @param x - Dividend array.
+ * @param y - Divisor array.
+ * @returns Element-wise floor division of x by y.
+ */
+declare function floorDivide(x: ArrayLike, y: ArrayLike): Array;
+/**
+ * @function
+ * Calculate element-wise floating-point modulo operation.
+ */
+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.
  *
- * @example
- * ```ts
- * import { lax, numpy as np } from "@jax-js/jax";
+ * Equivalent to `[floorDivide(x, y), remainder(x, y)]`.
  *
- * 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]
- * ```
+ * @param x - Dividend array.
+ * @param y - Divisor array.
+ * @returns Tuple of [quotient, remainder].
  */
-declare function lu(x: ArrayLike): [Array, Array, Array];
+declare function divmod(x: ArrayLike, y: ArrayLike): [Array, Array];
+/** Round input to the nearest integer towards zero. */
+declare function trunc(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.
+ * Compute `x1 * 2 ** x2` as a standard multiplication and exponentiation.
  *
- * @example
- * ```ts
- * import { lax, numpy as np } from "@jax-js/jax";
+ * This is the inverse of `frexp()`.
+ */
+declare function ldexp(x1: ArrayLike, x2: ArrayLike): Array;
+/**
+ * Decompose floating-point values into mantissa and two's exponent.
  *
- * const L = np.array([[2., 0.], [1., 3.]]);
- * const b = np.array([4., 7.]).reshape([2, 1]);
+ * 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 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;
+/**
+ * @function
+ * Computes first array raised to power of second array, element-wise.
+ */
+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>;
+/**
+ * @function
+ * Calculate element-wise hyperbolic sine of input.
  *
- * // Solve L @ x = b
- * const x = lax.linalg.triangularSolve(L, b, { leftSide: true, lower: true });
- * // x = [[2.], [5./3.]]
- * ```
+ * `sinh(x) = (exp(x) - exp(-x)) / 2`
  */
-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, erf, erfc, lax_linalg_d_exports as linalg, reduceWindow, stopGradient };
-}
+declare const sinh: OwnedFunction<(x: ArrayLike) => Array>;
 /**
- * Dimension numbers for general `dot()` primitive.
+ * @function
+ * Calculate element-wise 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.
+ * `cosh(x) = (exp(x) + exp(-x)) / 2`
+ */
+declare const cosh: OwnedFunction<(x: ArrayLike) => Array>;
+/**
+ * @function
+ * Calculate element-wise 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.
+ * `tanh(x) = sinh(x)/cosh(x) = (exp(x) - exp(-x)) / (exp(x) + exp(-x))`
  */
-type DotDimensionNumbers = {
-  lhsContractingDims?: number[];
-  rhsContractingDims?: number[];
-  lhsBatchDims?: number[];
-  rhsBatchDims?: number[];
-};
+declare const tanh: OwnedFunction<(x: ArrayLike) => Array>;
 /**
- * General dot product/contraction operator.
+ * @function
+ * Calculate element-wise inverse hyperbolic sine of input.
  *
- * Prefer higher-level functions like `jax.numpy.dot()`, `jax.numpy.matmul()`,
- * `jax.numpy.tensordot(), and `jax.numpy.einsum()` where possible.
+ * `arcsinh(x) = ln(x + sqrt(x^2 + 1))`
  */
-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 const arcsinh: OwnedFunction<(x: ArrayLike) => Array>;
 /**
- * General n-dimensional convolution operator, with optional dilation.
+ * @function
+ * Calculate element-wise inverse hyperbolic cosine of input.
  *
- * The semantics of this operation mimic the `jax.lax.conv_general_dilated`
- * function in JAX, which wraps XLA's general convolution operator.
+ * `arccosh(x) = ln(x + sqrt(x^2 - 1))`
+ */
+declare const arccosh: OwnedFunction<(x: ArrayLike) => Array>;
+/**
+ * @function
+ * Calculate element-wise inverse hyperbolic tangent of input.
  *
- * @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"`)
+ * `arctanh(x) = 0.5 * ln((1 + x) / (1 - x))`
  */
-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;
+declare const arctanh: OwnedFunction<(x: ArrayLike) => Array>;
 /**
- * Convenience wrapper for calculating the N-d convolution "transpose".
+ * Compute the variance of an array.
  *
- * This function directly calculates a fractionally strided conv rather than
- * indirectly calculating the gradient (transpose) of a forward convolution.
- * It is equivalent to the JAX version, except:
+ * The variance is computed for the flattened array by default, otherwise over
+ * the specified axis.
  *
- * - 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`.
+ * 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 var_(x: ArrayLike, axis?: Axis, opts?: {
+  mean?: ArrayLike;
+  correction?: number;
+} & ReduceOpts): Array;
+/**
+ * Compute the standard deviation of an array.
  *
- * 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.
+ * The standard deviation is computed for the flattened array by default,
+ * otherwise over the specified axis.
  *
- * @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]`
+ * 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 convTranspose(lhs: Array, rhs: Array, strides: number[], padding: PaddingType, {
-  rhsDilation,
-  transposeKernel
+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
 }?: {
-  rhsDilation?: number[];
-  transposeKernel?: boolean;
+  rowvar?: 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;
+/** 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, 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 };
 }