@jax-js/jax 0.0.3 → 0.0.5

package/dist/index.d.ts

@@ -123,6 +123,19 @@ declare class ShapeTracker {
 }
 //#endregion
 //#region src/utils.d.ts
+/**
+ * Set the debug level for verbose logging.
+ *
+ * 1. JIT compile logs
+ * 2. Shader code
+ * 3. Expressions and metadata
+ * 4. JIT programs, tuning details
+ * 5. Most verbose operation traces
+ *
+ * This is an experimental API and may change in behavior. Do not rely on this
+ * in production.
+ */
+declare function setDebug(level: number): void;
 /** @inline */
 type RecursiveArray<T> = T | RecursiveArray<T>[];
 interface FpHashable {
@@ -138,7 +151,7 @@ interface FpHashable {
 declare class FpHash {
   #private;
   value: bigint;
-  update(...values: (string | boolean | number | bigint | null | undefined | FpHashable)[]): this;
+  update(x: string | boolean | number | bigint | null | undefined | FpHashable): this;
   static hash(...values: (string | boolean | number | bigint | null | undefined | FpHashable)[]): bigint;
 }
 /** Run a function while caching it inline inside a `Map`. */
@@ -154,6 +167,31 @@ declare enum DType {
 }
 /** @inline */
 type DataArray = Float32Array<ArrayBuffer> | Int32Array<ArrayBuffer> | Uint32Array<ArrayBuffer> | Float16Array<ArrayBuffer>;
+/**
+ * Promote two dtypes to their join according to the type lattice.
+ *
+ * When performing operations between arrays of different types, we need to
+ * promote both operands to a common type that can represent values from both
+ * input types. This follows JAX's type promotion rules.
+ *
+ * **Type lattice:**
+ * ```text
+ * bool -> uint32 -> int32 -> float16 -> float32
+ *  weakType --^
+ * ```
+ *
+ * `weakType` represents weakly typed arrays. These are created for JS numbers,
+ * which default to float32 but "weak" so they cast to the dtype of any array
+ * they are first combined with, except `bool`.
+ *
+ * **Examples:**
+ * - `promoteTypes(bool, int32) → int32`
+ * - `promoteTypes(uint32, int32) → int32`
+ * - `promoteTypes(int32, float16) → float16`
+ * - `promoteTypes(float16, float32) → float32`
+ * - `promoteTypes(uint32, float32) → float32`
+ */
+declare function promoteTypes(dtype1: DType, dtype2: DType): DType;
 /**
  * Mathematical expression on scalar values.
  *
@@ -177,6 +215,8 @@ declare class AluExp implements FpHashable {
   static max(a: AluExp, b: AluExp): AluExp;
   static sin(a: AluExp): AluExp;
   static cos(a: AluExp): AluExp;
+  static asin(a: AluExp): AluExp;
+  static atan(a: AluExp): AluExp;
   static exp(a: AluExp): AluExp;
   static log(a: AluExp): AluExp;
   static sqrt(a: AluExp): AluExp;
@@ -262,6 +302,8 @@ declare enum AluOp {
   Max = "Max",
   Sin = "Sin",
   Cos = "Cos",
+  Asin = "Asin",
+  Atan = "Atan",
   Exp = "Exp",
   Log = "Log",
   Sqrt = "Sqrt",
@@ -354,8 +396,8 @@ declare class Reduction implements FpHashable {
 //#region src/backend.d.ts
 type Device = "cpu" | "wasm" | "webgpu";
 declare const devices: Device[];
-/** Set the default device backend (must be initialized). */
-declare function setDevice(device: Device): void;
+/** Configure the default device for arrays. */
+declare function defaultDevice(device?: Device): Device;
 /**
  * Initialize `jax-js` library backends.
  *
@@ -491,6 +533,8 @@ declare enum Primitive {
   RandomBits = "random_bits",
   Sin = "sin",
   Cos = "cos",
+  Asin = "asin",
+  Atan = "atan",
   Exp = "exp",
   Log = "log",
   Sqrt = "sqrt",
@@ -566,6 +610,7 @@ interface PrimitiveParamsImpl extends Record<Primitive, Record<string, any>> {
     outDim: number;
   };
   [Primitive.JitCall]: {
+    name: string;
     jaxpr: Jaxpr;
     numConsts: number;
   };
@@ -581,8 +626,10 @@ declare enum CompareOp {
   LessEqual = "less_equal",
 }
 /** @inline */
+type Axis = number | number[] | null;
+/** @inline */
 type ReduceOpts = {
-  keepDims?: boolean;
+  keepdims?: boolean;
 };
 type MainTrace = {
   level: number;
@@ -602,10 +649,40 @@ declare abstract class Trace {
   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;
 }
+/**
+ * Broadcast shapes and promote types with casting for two avals.
+ *
+ * This implements the weak type behavior described in `promoteTypes()`, but not
+ * implemented in that function as `weakType` is not passed.
+ */
+
 declare abstract class Tracer {
   /** @ignore */
   readonly _trace: Trace;
@@ -659,9 +736,20 @@ declare abstract class Tracer {
    * ```
    */
   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;
@@ -675,11 +763,11 @@ declare abstract class Tracer {
   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?: number | number[], opts?: ReduceOpts): this;
+  sum(axis?: Axis, opts?: ReduceOpts): this;
   /** Product of the array elements over a given axis. */
-  prod(axis?: number | number[], opts?: ReduceOpts): this;
+  prod(axis?: Axis, opts?: ReduceOpts): this;
   /** Compute the average of the array elements along the specified axis. */
-  mean(axis?: number | number[], opts?: ReduceOpts): this;
+  mean(axis?: Axis, opts?: ReduceOpts): this;
   /** Permute the dimensions of an array. Defaults to reversing the axis order. */
   transpose(perm?: number[]): this;
   /**
@@ -752,7 +840,8 @@ declare abstract class Tracer {
 declare class ShapedArray implements AbstractValue {
   readonly shape: number[];
   readonly dtype: DType;
-  constructor(shape: number[], dtype: DType);
+  readonly weakType: boolean;
+  constructor(shape: number[], dtype: DType, weakType: boolean);
   static fromAval(aval: AbstractValue): ShapedArray;
   get ndim(): number;
   toString(): string;
@@ -788,6 +877,14 @@ type DTypeAndDevice = {
   dtype?: DType;
   device?: Device;
 };
+type ArrayConstructorArgs = {
+  source: AluExp | Slot;
+  st: ShapeTracker;
+  dtype: DType;
+  weakType: boolean;
+  backend: Backend;
+  pending?: Iterable<PendingExecute>;
+};
 /**
  * A multidimensional numeric array with data stored on CPU or GPU.
  *
@@ -807,7 +904,7 @@ declare class Array extends Tracer {
    * is a backend `Slot`, this constructor _takes ownership_ of the slot. It
    * will be freed when the array is disposed.
    */
-  constructor(source: AluExp | Slot, st: ShapeTracker, dtype: DType, backend: Backend, pending?: Iterable<PendingExecute> | null);
+  constructor(args: ArrayConstructorArgs);
   /** @ignore */
   get aval(): ShapedArray;
   /** Return a simple string representation of the array's dimensions. */
@@ -836,8 +933,11 @@ declare class Array extends Tracer {
    *
    * 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.
    */
-  wait(): Promise<Array>;
+  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.
@@ -866,11 +966,6 @@ declare class Array extends Tracer {
   static _implRules(): typeof implRules;
   _realizeSource(): number;
 }
-/** Construct an array from a single scalar constant. */
-declare function scalar(value: number | boolean, {
-  dtype,
-  device
-}?: DTypeAndDevice): Array;
 /** Constructor for creating a new array from data. */
 declare function array(values: Array | Float16Array<ArrayBuffer> | Float32Array<ArrayBuffer> | Int32Array<ArrayBuffer> | Uint32Array<ArrayBuffer> | RecursiveArray<number> | RecursiveArray<boolean>, {
   shape,
@@ -908,7 +1003,7 @@ declare function eye(numRows: number, numCols?: number, {
   dtype,
   device
 }?: DTypeAndDevice): Array;
-/** Return the identity array, with ones on the main diagonal. */
+/** Return the identity matrix, with ones on the main diagonal. */
 declare function identity$1(n: number, {
   dtype,
   device
@@ -945,7 +1040,7 @@ declare function linspace(start: number, stop: number, num?: number, endpoint?:
   device
 }?: DTypeAndDevice): Array;
 declare namespace numpy_d_exports {
-  export { Array, ArrayLike, DType, abs, absolute, add, allclose, arange, argmax, argmin, array, astype, bool, clip, columnStack, concatenate, cos, cosh, diag, diagonal, divide, dot, dstack, e, equal, eulerGamma, exp, exp2, eye, flip, fliplr, flipud, float16, float32, full, fullLike, greater, greaterEqual, hstack, identity$1 as identity, inf, int32, less, lessEqual, linspace, log, log10, log2, matmul, max, maximum, mean, meshgrid, min, minimum, moveaxis, multiply, nan, ndim, negative, notEqual, ones, onesLike, pad, permuteDims, pi, prod, ravel, reciprocal, reshape, scalar, shape$1 as shape, sin, sinh, size, sqrt, square, stack, sum, tan, tanh, transpose, trueDivide, trunc, uint32, vdot, vecdot, vstack, where, zeros, zerosLike };
+  export { Array, ArrayLike, DType, abs, absolute, acos, acosh, add, allclose, arange, arccos, arccosh, arcsinh, arctan, arctan2, arctanh, argmax, argmin, array, asin, asinh, astype, atan, atan2, atanh, bool, broadcastArrays, broadcastShapes, broadcastTo, cbrt, clip, columnStack, concatenate, cos, cosh, deg2rad, degrees, diag, diagonal, divide, dot, dstack, e, equal, eulerGamma, exp, exp2, expm1, eye, flip, fliplr, flipud, float16, float32, full, fullLike, greater, greaterEqual, hstack, hypot, identity$1 as identity, inf, inner, int32, less, lessEqual, linspace, log, log10, log1p, log2, matmul, max, maximum, mean, meshgrid, min, minimum, moveaxis, multiply, nan, ndim, negative, notEqual, ones, onesLike, outer, pad, permuteDims, pi, pow, power, prod, promoteTypes, rad2deg, radians, ravel, reciprocal, repeat, reshape, shape$1 as shape, sign, sin, sinh, size, sqrt, square, stack, std, subtract, sum, tan, tanh, tile, transpose, tri, tril, triu, trueDivide, trunc, uint32, var_, vdot, vecdot, vstack, where, zeros, zerosLike };
 }
 declare const float32 = DType.Float32;
 declare const int32 = DType.Int32;
@@ -962,54 +1057,66 @@ declare const inf: number;
 declare const nan: number;
 /** This is Pi, `π = 3.14159265358979...` */
 declare const pi: number;
-/** Element-wise addition, with broadcasting. */
+/** @function Element-wise addition, with broadcasting. */
 declare const add: (x: ArrayLike, y: ArrayLike) => Array;
-/** Element-wise multiplication, with broadcasting. */
+/** @function Element-wise multiplication, with broadcasting. */
 declare const multiply: (x: ArrayLike, y: ArrayLike) => Array;
-/** Numerical negative of every element of an array. */
+/** @function Numerical negative of every element of an array. */
 declare const negative: (x: ArrayLike) => Array;
-/** Calculate element-wise reciprocal of the input. This is `1/x`. */
+/** @function Calculate element-wise reciprocal of the input. This is `1/x`. */
 declare const reciprocal: (x: ArrayLike) => Array;
-/** Element-wise sine function (takes radians). */
+/** @function Element-wise sine function (takes radians). */
 declare const sin: (x: ArrayLike) => Array;
-/** Element-wise cosine function (takes radians). */
+/** @function Element-wise cosine function (takes radians). */
 declare const cos: (x: ArrayLike) => Array;
-/** Calculate the exponential of all elements in the input 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;
-/** Calculate the natural logarithm of all elements in the input array. */
+/** @function Calculate the natural logarithm of all elements in the input array. */
 declare const log: (x: ArrayLike) => Array;
-/** Calculate the square root of all elements in the input array. */
+/** @function Calculate the square root of all elements in the input array. */
 declare const sqrt: (x: ArrayLike) => Array;
-/** Return element-wise minimum of the input arrays. */
+/** @function Return element-wise minimum of the input arrays. */
 declare const minimum: (x: ArrayLike, y: ArrayLike) => Array;
-/** Return element-wise maximum of the input arrays. */
+/** @function Return element-wise maximum of the input arrays. */
 declare const maximum: (x: ArrayLike, y: ArrayLike) => Array;
-/** Compare two arrays element-wise. */
+/** @function Compare two arrays element-wise. */
 declare const greater: (x: ArrayLike, y: ArrayLike) => Array;
-/** Compare two arrays element-wise. */
+/** @function Compare two arrays element-wise. */
 declare const less: (x: ArrayLike, y: ArrayLike) => Array;
-/** Compare two arrays element-wise. */
+/** @function Compare two arrays element-wise. */
 declare const equal: (x: ArrayLike, y: ArrayLike) => Array;
-/** Compare two arrays element-wise. */
+/** @function Compare two arrays element-wise. */
 declare const notEqual: (x: ArrayLike, y: ArrayLike) => Array;
-/** Compare two arrays element-wise. */
+/** @function Compare two arrays element-wise. */
 declare const greaterEqual: (x: ArrayLike, y: ArrayLike) => Array;
-/** Compare two arrays element-wise. */
+/** @function Compare two arrays element-wise. */
 declare const lessEqual: (x: ArrayLike, y: ArrayLike) => Array;
-/** Element-wise ternary operator, evaluates to `x` if cond else `y`. */
+/** @function Element-wise ternary operator, evaluates to `x` if cond else `y`. */
 declare const where: (cond: ArrayLike, x: ArrayLike, y: ArrayLike) => Array;
-/** Permute the dimensions of an array. Defaults to reversing the axis order. */
+/**
+ * @function
+ * Permute the dimensions of an array. Defaults to reversing the axis order.
+ */
 declare const transpose: (x: ArrayLike, perm?: number[]) => Array;
 /**
+ * @function
  * 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.
  */
 declare const reshape: (x: ArrayLike, shape: number[]) => Array;
-/** Move axes of an array to new positions. Other axes retain original order. */
+/**
+ * @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
@@ -1017,15 +1124,27 @@ declare const moveaxis: (x: ArrayLike, src: number, dst: number) => Array;
  * pair specifies the padding for its corresponding axis.
  */
 declare const pad: (x: ArrayLike, width: number | [number, number] | [number, number][]) => Array;
-/** Return the number of dimensions of an array. Does not consume array reference. */
+/**
+ * @function
+ * Return the number of dimensions of an array. Does not consume array reference.
+ */
 declare const ndim: (x: ArrayLike) => number;
-/** Return the shape of an array. Does not consume array reference. */
+/** @function Return the shape of an array. Does not consume array reference. */
 declare const shape$1: (x: ArrayLike) => number[];
-/** Return an array of zeros with the same shape and type as a given array. */
+/**
+ * @function
+ * Return an array of zeros with the same shape and type as a given array.
+ */
 declare const zerosLike: (a: ArrayLike, dtype?: DType) => Array;
-/** Return an array of ones with the same shape and type as a given 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;
-/** Return a full array with the same shape and type as a given 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.
@@ -1035,15 +1154,15 @@ 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?: number | number[], opts?: ReduceOpts): Array;
+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?: number | number[], opts?: ReduceOpts): Array;
+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?: number | number[], opts?: ReduceOpts): Array;
+declare function min(a: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
 /** Return the maximum of array elements along a given axis. */
-declare function max(a: ArrayLike, axis?: number | number[], opts?: ReduceOpts): Array;
+declare function max(a: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
 /** Compute the average of the array elements along the specified axis. */
-declare function mean(a: ArrayLike, axis?: number | number[], opts?: ReduceOpts): Array;
+declare function mean(a: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
 /**
  * Returns the indices of the minimum values along an axis.
  *
@@ -1059,7 +1178,7 @@ declare function argmin(a: ArrayLike, axis?: number, opts?: ReduceOpts): Array;
  */
 declare function argmax(a: ArrayLike, axis?: number, opts?: ReduceOpts): Array;
 /** Reverse the elements in an array along the given axes. */
-declare function flip(x: ArrayLike, axis?: number | number[]): Array;
+declare function flip(x: ArrayLike, axis?: Axis): Array;
 /**
  * Join a sequence of arrays along an existing axis.
  *
@@ -1103,9 +1222,36 @@ declare function columnStack(xs: ArrayLike[]): Array;
 declare function flipud(x: ArrayLike): Array;
 /** Flip an array horizontally (axis=1). */
 declare function fliplr(x: ArrayLike): Array;
+/** @function Alternative name for `numpy.transpose()`. */
 declare const permuteDims: (x: ArrayLike, perm?: number[]) => Array;
 /** Return a 1-D flattened array containing the elements of the input. */
 declare function ravel(a: ArrayLike): 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.
  *
@@ -1133,8 +1279,28 @@ declare function allclose(actual: Parameters<typeof array>[0], expected: Paramet
 declare function matmul(x: ArrayLike, y: ArrayLike): Array;
 /** Dot product of two arrays. */
 declare function dot(x: ArrayLike, y: ArrayLike): Array;
-/** Vector dot product of two arrays. */
-declare function vecdot(x: ArrayLike, y: ArrayLike): Array;
+/**
+ * Compute the inner product of two arrays.
+ *
+ * Unlike `jax.numpy.matmul()` or `jax.numpy.dot()`, this always performs a
+ * contraction on the last axis.
+ *
+ * Returned array has shape `[...x.shape[:-1], ...y.shape[:-1]]`.
+ */
+declare function inner(x: ArrayLike, y: ArrayLike): Array;
+/**
+ * Compute the outer product of two arrays.
+ *
+ * If the input arrays are not 1D, they will be flattened. Returned array will
+ * be of shape `[x.size, y.size]`.
+ */
+declare function outer(x: ArrayLike, y: ArrayLike): Array;
+/** Vector dot product of two arrays along a given axis. */
+declare function vecdot(x: ArrayLike, y: ArrayLike, {
+  axis
+}?: {
+  axis?: number;
+}): Array;
 /**
  * Return the dot product of two vectors.
  *
@@ -1152,6 +1318,21 @@ declare function meshgrid(xs: Array[], {
 }?: {
   indexing?: "xy" | "ij";
 }): Array[];
+/**
+ * Return an array with ones on and below the diagonal and zeros elsewhere.
+ *
+ * If `k` is provided, it specifies the sub-diagonal on and below which the
+ * array is filled with ones. `k=0` is the main diagonal, `k<0` is below it, and
+ * `k>0` is above it.
+ */
+declare function tri(n: number, m?: number, k?: number, {
+  dtype,
+  device
+}?: DTypeAndDevice): Array;
+/** Return the lower triangle of an array. Must be of dimension >= 2. */
+declare function tril(a: ArrayLike, k?: number): Array;
+/** Return the upper triangle of an array. Must be of dimension >= 2. */
+declare function triu(a: ArrayLike, k?: number): Array;
 /**
  * Clip (limit) the values in an array.
  *
@@ -1168,15 +1349,50 @@ declare function clip(a: ArrayLike, min?: ArrayLike, max?: ArrayLike): Array;
  * This is the same function as `jax.numpy.abs()`.
  */
 declare function absolute(x: ArrayLike): Array;
-/** Alias of `jax.numpy.absolute()`. */
+/** @function Alias of `jax.numpy.absolute()`. */
 declare const abs: typeof absolute;
+/** Return an element-wise indication of sign of the input. */
+declare function sign(x: ArrayLike): Array;
 /** Calculate element-wise square of the input array. */
 declare function square(x: ArrayLike): Array;
-/** Compute a trigonometric tangent of each element of input. */
+/** Element-wise tangent function (takes radians). */
 declare function tan(x: ArrayLike): Array;
+/** Element-wise inverse cosine function (inverse of cos). */
+declare function acos(x: ArrayLike): Array;
+/**
+ * @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>;
+/** @function Alias of `jax.numpy.acos()`. */
+declare const arccos: typeof acos;
+/** @function Alias of `jax.numpy.atan()`. */
+declare const arctan: (x: ArrayLike) => Array;
+/** @function Alias of `jax.numpy.atan2()`. */
+declare const arctan2: OwnedFunction<(y: ArrayLike, x: ArrayLike) => Array>;
+/** Element-wise subtraction, with broadcasting. */
+declare function subtract(x: ArrayLike, y: ArrayLike): Array;
 /** Calculates the floating-point division of x by y element-wise. */
 declare function trueDivide(x: ArrayLike, y: ArrayLike): Array;
-/** Alias of `jax.numpy.trueDivide()`. */
+/** @function Alias of `jax.numpy.trueDivide()`. */
 declare const divide: typeof trueDivide;
 /** Round input to the nearest integer towards zero. */
 declare function trunc(x: ArrayLike): Array;
@@ -1186,26 +1402,112 @@ declare function exp2(p: ArrayLike): Array;
 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 Alias of `jax.numpy.power()`. */
+declare const pow: OwnedFunction<(x1: ArrayLike, x2: ArrayLike) => Array>;
+/** @function Calculate the element-wise cube root of the input array. */
+declare const cbrt: OwnedFunction<(x: ArrayLike) => Array>;
 /**
+ * @function
  * Calculate element-wise hyperbolic sine of input.
  *
  * `sinh(x) = (exp(x) - exp(-x)) / 2`
  */
-declare function sinh(x: ArrayLike): Array;
+declare const sinh: OwnedFunction<(x: ArrayLike) => Array>;
 /**
+ * @function
  * Calculate element-wise hyperbolic cosine of input.
  *
  * `cosh(x) = (exp(x) + exp(-x)) / 2`
  */
-declare function cosh(x: ArrayLike): Array;
+declare const cosh: OwnedFunction<(x: ArrayLike) => Array>;
 /**
+ * @function
  * Calculate element-wise hyperbolic tangent of input.
  *
  * `tanh(x) = sinh(x)/cosh(x) = (exp(x) - exp(-x)) / (exp(x) + exp(-x))`
  */
-declare function tanh(x: ArrayLike): Array;
+declare const tanh: OwnedFunction<(x: ArrayLike) => Array>;
+/**
+ * @function
+ * Calculate element-wise inverse hyperbolic sine of input.
+ *
+ * `arcsinh(x) = ln(x + sqrt(x^2 + 1))`
+ */
+declare const arcsinh: OwnedFunction<(x: ArrayLike) => Array>;
+/**
+ * @function
+ * Calculate element-wise inverse hyperbolic cosine of input.
+ *
+ * `arccosh(x) = ln(x + sqrt(x^2 - 1))`
+ */
+declare const arccosh: OwnedFunction<(x: ArrayLike) => 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>;
+/** @function Alias of `jax.numpy.arcsinh()`. */
+declare const asinh: OwnedFunction<(x: ArrayLike) => Array>;
+/** @function Alias of `jax.numpy.arccosh()`. */
+declare const acosh: OwnedFunction<(x: ArrayLike) => Array>;
+/** @function Alias of `jax.numpy.arctanh()`. */
+declare const atanh: OwnedFunction<(x: ArrayLike) => Array>;
+/**
+ * Compute the variance of an array.
+ *
+ * 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).
+ */
+declare function std(x: ArrayLike, axis?: Axis, opts?: {
+  mean?: ArrayLike;
+  correction?: number;
+} & ReduceOpts): Array;
 //#endregion
 //#region src/frontend/jaxpr.d.ts
+/**
+ * Function callback with an associated dispose() method.
+ *
+ * The dispose() method should be called to clean up any tracer resources needed
+ * by the function after the last time it is called.
+ */
+type OwnedFunction<F extends Function> = F & {
+  dispose: () => void;
+};
 /** Variable in a Jaxpr expression. */
 declare class Var {
   #private;
@@ -1216,10 +1518,10 @@ declare class Var {
 }
 /** Literal in a Jaxpr expression. Currently, only scalars are supported. */
 declare class Lit {
-  readonly dtype: DType;
   readonly value: number;
   readonly aval: ShapedArray;
-  constructor(dtype: DType, value: number);
+  get dtype(): DType;
+  constructor(aval: AbstractValue, value: number);
 }
 type Atom = Var | Lit;
 declare class VarPrinter {
@@ -1297,7 +1599,7 @@ declare function conv(lhs: Array, rhs: Array, windowStrides: number[], padding:
 /** Reduce a computation over padded windows. */
 declare function reduceWindow(operand: Array, computation: (x: Array) => Array, windowDimensions: number[], windowStrides?: number[]): Array;
 declare namespace nn_d_exports {
-  export { celu, elu, gelu, glu, identity, leakyRelu, logSigmoid, logSoftmax, logsumexp, mish, oneHot, relu, relu6, sigmoid, silu, softSign, softmax, softplus, swish };
+  export { celu, elu, gelu, glu, identity, leakyRelu, logSigmoid, logSoftmax, logmeanexp, logsumexp, mish, oneHot, relu, relu6, sigmoid, silu, softSign, softmax, softplus, squareplus, standardize, swish };
 }
 /**
  * Rectified Linear Unit (ReLU) activation function:
@@ -1329,6 +1631,7 @@ declare function softplus(x: ArrayLike): Array;
  */
 declare function softSign(x: ArrayLike): Array;
 /**
+ * @function
  * Sigmoid-weighted Linear Unit (SiLU) activation function, also known as
  * Swish, computed element-wise:
  * `silu(x) = x * sigmoid(x) = x / (1 + exp(-x))`.
@@ -1337,8 +1640,9 @@ declare function softSign(x: ArrayLike): Array;
  *
  * Reference: https://en.wikipedia.org/wiki/Swish_function
  */
-declare const silu: (x: ArrayLike) => Array;
+declare const silu: OwnedFunction<(x: ArrayLike) => Array>;
 /**
+ * @function
  * Sigmoid-weighted Linear Unit (SiLU) activation function, also known as
  * Swish, computed element-wise:
  * `silu(x) = x * sigmoid(x) = x / (1 + exp(-x))`.
@@ -1347,31 +1651,35 @@ declare const silu: (x: ArrayLike) => Array;
  *
  * Reference: https://en.wikipedia.org/wiki/Swish_function
  */
-declare const swish: (x: ArrayLike) => Array;
+declare const swish: OwnedFunction<(x: ArrayLike) => Array>;
 /**
  * Log-sigmoid activation function, computed element-wise:
  * `log_sigmoid(x) = log(sigmoid(x)) = -log(1 + exp(-x))`.
  */
 declare function logSigmoid(x: ArrayLike): Array;
-/** Identity activation function. Returns the argument unmodified. */
+/**
+ * @function
+ * Identity activation function. Returns the argument unmodified.
+ */
 declare const identity: (x: ArrayLike) => Array;
 /** Leaky rectified linear (ReLU) activation function */
-declare function leakyRelu(x: ArrayLike, negativeSlope?: number): Array;
+declare function leakyRelu(x: ArrayLike, negativeSlope?: ArrayLike): Array;
 /**
  * Exponential linear unit activation function.
  *
  * Computes the element-wise function:
  * `elu(x) = x > 0 ? x : alpha * (exp(x) - 1)`
  */
-declare function elu(x: ArrayLike, alpha?: number): Array;
+declare function elu(x: ArrayLike, alpha?: ArrayLike): Array;
 /**
  * Continuously-differentiable exponential linear unit activation function.
  *
  * Computes the element-wise function:
  * `celu(x) = x > 0 ? x : alpha * (exp(x/alpha) - 1)`
  */
-declare function celu(x: ArrayLike, alpha?: number): Array;
+declare function celu(x: ArrayLike, alpha?: ArrayLike): Array;
 /**
+ * @function
  * Gaussion error linear unit (GELU) activation function.
  *
  * This is computed element-wise. Currently jax-js does not support the erf() or
@@ -1382,7 +1690,7 @@ declare function celu(x: ArrayLike, alpha?: number): Array;
  *
  * This will be improved in the future.
  */
-declare const gelu: (x: ArrayLike) => Array;
+declare const gelu: OwnedFunction<(x: ArrayLike) => Array>;
 /**
  * Gated linear unit (GLU) activation function.
  *
@@ -1390,6 +1698,13 @@ declare const gelu: (x: ArrayLike) => Array;
  * computes `a * sigmoid(b)`.
  */
 declare function glu(x: ArrayLike, axis?: number): Array;
+/**
+ * Squareplus activation function.
+ *
+ * Computes the element-wise function:
+ * `squareplus(x) = 0.5 * (x + sqrt(x^2 + b))`
+ */
+declare function squareplus(x: ArrayLike, b?: ArrayLike): Array;
 /**
  * Mish activation function.
  *
@@ -1405,7 +1720,7 @@ declare function mish(x: ArrayLike): Array;
  *
  * Reference: https://en.wikipedia.org/wiki/Softmax_function
  */
-declare function softmax(x: ArrayLike, axis?: number | number[]): Array;
+declare function softmax(x: ArrayLike, axis?: Axis): Array;
 /**
  * Log-Softmax function.
  *
@@ -1414,7 +1729,7 @@ declare function softmax(x: ArrayLike, axis?: number | number[]): Array;
  *
  * If `axis` is not specified, it defaults to the last axis.
  */
-declare function logSoftmax(x: ArrayLike, axis?: number | number[]): Array;
+declare function logSoftmax(x: ArrayLike, axis?: Axis): Array;
 /**
  * Log-sum-exp reduction. Also a multivariate version of `softplus`.
  *
@@ -1423,7 +1738,22 @@ declare function logSoftmax(x: ArrayLike, axis?: number | number[]): Array;
  *
  * Reference: https://en.wikipedia.org/wiki/LogSumExp
  */
-declare function logsumexp(x: ArrayLike, axis?: number | number[]): Array;
+declare function logsumexp(x: ArrayLike, axis?: Axis): Array;
+/** Log-mean-exp reduction, like `jax.nn.logsumexp()` but subtracts `log(n)`. */
+declare function logmeanexp(x: ArrayLike, axis?: Axis): Array;
+/**
+ * Standardizes input to zero mean and unit variance.
+ *
+ * By default, this is computed over the last axis. You can pass in a different
+ * axis, or `null` to standardize over all elements.
+ *
+ * Epsilon is added to denominator, it defaults to `1e-5` for stability.
+ */
+declare function standardize(x: ArrayLike, axis?: Axis, opts?: {
+  mean?: ArrayLike;
+  variance?: ArrayLike;
+  epsilon?: ArrayLike;
+}): Array;
 /**
  * One-hot encodes the given indices.
  *
@@ -1442,7 +1772,7 @@ declare function logsumexp(x: ArrayLike, axis?: number | number[]): Array;
  */
 declare function oneHot(x: Array, numClasses: number): Array;
 declare namespace random_d_exports {
-  export { bits, key, split, uniform };
+  export { bernoulli, bits, exponential, key, normal, split, uniform };
 }
 /** Create a pseudo-random number generator (PRNG) key from 32-bit integer seed. */
 declare function key(seed: number): Array;
@@ -1450,34 +1780,71 @@ declare function key(seed: number): Array;
 declare function split(key: Array, num?: number | number[]): Array;
 /** Sample uniform bits in the form of unsigned integers. */
 declare function bits(key: Array, shape?: number[]): Array;
-/** Sample uniform random values in [minval, maxval) with given shape. */
-declare function uniform(key: Array, shape?: number[], {
-  minval,
-  maxval
-}?: {
-  minval?: number;
-  maxval?: number;
-}): Array;
+/**
+ * @function
+ * Sample uniform random values in [minval, maxval) with given shape.
+ */
+declare const uniform: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined, args_2?: {
+  minval?: number | undefined;
+  maxval?: number | undefined;
+} | undefined) => Array>;
+/**
+ * Sample Bernoulli random variables with given mean (0,1 categorical).
+ *
+ * Returns a random Boolean array with the specified shape. `p` can be an array
+ * and must be broadcastable to `shape`.
+ */
+declare function bernoulli(key: Array, p?: ArrayLike, shape?: number[]): Array;
+/**
+ * @function
+ * Sample exponential random values according to `p(x) = exp(-x)`.
+ */
+declare const exponential: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined) => Array>;
+/**
+ * @function
+ * Sample random values according to `p(x) = 1/sqrt(2pi) * exp(-x^2/2)`.
+ *
+ * Unlike JAX, this uses the Box-Muller transform. JAX uses the erf_inv primitive instead and
+ * directly inverts the CDF, but we don't have support for that yet. Outputs will not be
+ * bitwise identical to JAX.
+ */
+declare const normal: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined) => Array>;
 //#endregion
 //#region src/index.d.ts
-/** Compute the forward-mode Jacobian-vector product for a function. */
+/**
+ * @function
+ * Compute the forward-mode Jacobian-vector product for a function.
+ */
 declare const jvp: <F extends (...args: any[]) => JsTree<Array>>(f: F, primals: MapJsTree<Parameters<F>, Array, ArrayLike>, tangents: MapJsTree<Parameters<F>, Array, ArrayLike>) => [ReturnType<F>, ReturnType<F>];
-/** Vectorize an operation on a batched axis for one or more inputs. */
+/**
+ * @function
+ * Vectorize an operation on a batched axis for one or more inputs.
+ */
 declare const vmap: <F extends (...args: any[]) => JsTree<Array>>(f: F, inAxes?: number | MapJsTree<Parameters<F>, ArrayLike, number | null>) => (...args: MapJsTree<Parameters<F>, Array, ArrayLike>) => ReturnType<F>;
-/** Compute the Jacobian evaluated column-by-column by forward-mode AD. */
+/**
+ * @function
+ * Compute the Jacobian evaluated column-by-column by forward-mode AD.
+ */
 declare const jacfwd: <F extends (x: Array) => Array>(f: F) => (...args: MapJsTree<Parameters<F>, Array, ArrayLike>) => ReturnType<F>;
-/** Construct a Jaxpr by dynamically tracing a function with example inputs. */
+/**
+ * @function
+ * Construct a Jaxpr by dynamically tracing a function with example inputs.
+ */
 declare const makeJaxpr: <F extends (...args: any[]) => JsTree<Array>>(f: F) => (...args: Parameters<F>) => {
   jaxpr: Jaxpr;
   consts: Array[];
   treedef: JsTreeDef;
 };
 /**
+ * @function
  * Mark a function for automatic JIT compilation, with operator fusion.
  *
  * The function will be compiled the first time it is called with a set of
  * argument shapes.
  *
+ * You can call `.dispose()` on the returned, JIT-compiled function after all
+ * calls to free memory associated with array constants.
+ *
  * **Options:**
  * - `staticArgnums`: An array of argument indices to treat as static
  *   (compile-time constant). These arguments must be hashable, won't be traced,
@@ -1485,24 +1852,48 @@ declare const makeJaxpr: <F extends (...args: any[]) => JsTree<Array>>(f: F) =>
  * - `device`: The device to place the computation on. If not specified, the
  *   computation will be placed on the default device.
  */
-declare const jit: <F extends (...args: any[]) => JsTree<Array>>(f: F, opts?: JitOpts) => (...args: MapJsTree<Parameters<F>, Array, ArrayLike>) => ReturnType<F>;
+declare const jit: <F extends (...args: any[]) => JsTree<Array>>(f: F, opts?: JitOpts) => OwnedFunction<(...args: MapJsTree<Parameters<F>, Array, ArrayLike>) => ReturnType<F>>;
 /**
+ * @function
  * Produce a local linear approximation to a function at a point using jvp() and
  * partial evaluation.
  */
 declare const linearize: <F extends (...args: any[]) => JsTree<Array>>(f: F, ...primals: MapJsTree<Parameters<F>, Array, ArrayLike>) => [ReturnType<F>, (...tangents: MapJsTree<Parameters<F>, Array, ArrayLike>) => ReturnType<F>];
-/** Calculate the reverse-mode vector-Jacobian product for a function. */
+/**
+ * @function
+ * Calculate the reverse-mode vector-Jacobian product for a function.
+ */
 declare const vjp: <F extends (...args: any[]) => JsTree<Array>>(f: F, ...primals: MapJsTree<Parameters<F>, Array, ArrayLike>) => [ReturnType<F>, (cotangents: MapJsTree<ReturnType<F>, Array, ArrayLike>) => MapJsTree<Parameters<F>, ArrayLike, Array>];
 /**
+ * @function
  * Compute the gradient of a scalar-valued function `f` with respect to its
  * first argument.
  */
 declare const grad: <F extends (...args: any[]) => JsTree<Array>>(f: F) => (...primals: MapJsTree<Parameters<F>, Array, ArrayLike>) => MapJsTree<Parameters<F>[0], ArrayLike, Array>;
-/** Create a function that evaluates both `f` and the gradient of `f`. */
+/**
+ * @function
+ * Create a function that evaluates both `f` and the gradient of `f`.
+ */
 declare const valueAndGrad: <F extends (...args: any[]) => JsTree<Array>>(f: F) => (...primals: MapJsTree<Parameters<F>, Array, ArrayLike>) => [ReturnType<F>, MapJsTree<Parameters<F>[0], ArrayLike, Array>];
-/** Compute the Jacobian evaluated row-by-row by reverse-mode AD. */
+/**
+ * @function
+ * Compute the Jacobian evaluated row-by-row by reverse-mode AD.
+ */
 declare const jacrev: typeof jacfwd;
-/** Compute the Jacobian with reverse-mode AD. Alias for `jacrev()`. */
+/**
+ * @function
+ * Compute the Jacobian with reverse-mode AD. Alias for `jacrev()`.
+ */
 declare const jacobian: <F extends (x: Array) => Array>(f: F) => (...args: MapJsTree<Parameters<F>, Array, ArrayLike>) => ReturnType<F>;
+/**
+ * Wait until all `Array` leaves are ready by calling `Array.blockUntilReady()`.
+ *
+ * This can be used to wait for the results of an intermediate computation to
+ * finish. It's recommended to call this regularly in an iterative computation
+ * to avoid queueing up too many pending operations.
+ *
+ * Does not consume reference to the arrays.
+ */
+declare function blockUntilReady<T extends JsTree<any>>(x: T): Promise<T>;
 //#endregion
-export { DType, type Device, type JsTree, type JsTreeDef, devices, grad, init, jacfwd, jacobian, jacrev, jit, jvp, lax_d_exports as lax, linearize, makeJaxpr, nn_d_exports as nn, numpy_d_exports as numpy, random_d_exports as random, setDevice, tree_d_exports as tree, valueAndGrad, vjp, vmap };
+export { Array, DType, type Device, Jaxpr, type JsTree, type JsTreeDef, type OwnedFunction, blockUntilReady, defaultDevice, devices, grad, init, jacfwd, jacobian, jacrev, jit, jvp, lax_d_exports as lax, linearize, makeJaxpr, nn_d_exports as nn, numpy_d_exports as numpy, random_d_exports as random, setDebug, tree_d_exports as tree, valueAndGrad, vjp, vmap };