@jax-js/jax 0.0.2 → 0.0.4

package/dist/index.d.ts

@@ -20,6 +20,7 @@ interface Stringable {
 //#endregion
 //#region src/shape.d.ts
 
+/** @inline */
 type Pair = [number, number];
 /**
  * A multidimensional view into memory. An array can be thought of as the
@@ -46,6 +47,8 @@ declare class View {
   get size(): number;
   /** Whether this is a default, contiguous, unaltered view of the data (identity). */
   get contiguous(): boolean;
+  /** Return the range of data being indexed in this view, or [0, 0] if none. */
+  dataRange(): [number, number];
   /** Produce an AluExp for evaluating this view at an index. */
   toAluExp(idxs: AluExp[]): [AluExp, AluExp];
   /**
@@ -106,9 +109,33 @@ declare class ShapeTracker {
   reshape(newShape: number[]): ShapeTracker;
   /** Broadcast along the given new axes, then expand the shape. */
   broadcast(newShape: number[], axis: number[]): ShapeTracker;
+  /**
+   * Repeat data in each axis by a positive number of repetitions.
+   *
+   * - If `tile` is true (default): [1, 2, 3] -> [1, 2, 3, 1, 2, 3].
+   * - If `tile` is false: [1, 2, 3] -> [1, 1, 2, 2, 3, 3].
+   */
+  repeat(reps: number[], tile?: boolean): ShapeTracker;
+  /** Move axis i to axis j. */
+  moveaxis(i: number, j: number): ShapeTracker;
+  /** Like pad(), but allows for negative values. */
+  padOrShrink(arg: Pair[]): 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 {
@@ -124,19 +151,47 @@ 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`. */
 //#endregion
 //#region src/alu.d.ts
+/** A numerical data type for array contents. */
 declare enum DType {
   Float32 = "float32",
   Int32 = "int32",
   Uint32 = "uint32",
   Bool = "bool",
-  Complex64 = "complex64",
+  Float16 = "float16",
 }
+/** @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
+ *  weak f* --^
+ * ```
+ *
+ * The asterisk f* is a weak type used for JS number constants. When creating
+ * arrays, JS numbers default to float32 but "weak" so they cast to the dtype of
+ * any array they are first combined with.
+ *
+ * **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.
  *
@@ -160,8 +215,11 @@ 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;
   static reciprocal(a: AluExp): AluExp;
   static cast(dtype: DType, a: AluExp): AluExp;
   static bitcast(dtype: DType, a: AluExp): AluExp;
@@ -172,12 +230,13 @@ declare class AluExp implements FpHashable {
   static const(dtype: DType, value: any): AluExp;
   static special(dtype: DType, name: string, n: number): AluExp;
   static variable(dtype: DType, name: string): AluExp;
-  static globalIndex(dtype: DType, gid: number, bufidx: AluExp): AluExp;
+  static globalIndex(dtype: DType, gid: number, len: number, bufidx: AluExp): AluExp;
   static globalView(dtype: DType, gid: number, st: ShapeTracker, indices: AluExp[]): AluExp;
+  static f32(value: number): AluExp;
   static i32(value: number): AluExp;
   static u32(value: number): AluExp;
-  static f32(value: number): AluExp;
   static bool(value: boolean): AluExp;
+  static f16(value: number): AluExp;
   not(): AluExp;
   /** Compute a reasonable expression hash with low collision rate. */
   getHash(): bigint;
@@ -188,6 +247,19 @@ declare class AluExp implements FpHashable {
   reindexGids(gidMap: Map<number, number>): AluExp;
   get min(): number;
   get max(): number;
+  /** Largest known integer that divides self. */
+  constFactor(): number;
+  /**
+   * Checks if divisible by an integer v and returns the quotient if it is, or
+   * `null` if it's not divisible.
+   */
+  divides(v: number): AluExp | null;
+  /**
+   * Get all expressions by deeply matching an operation.
+   *
+   * For example: `((2+(3*5))+4).splitOp(+) -> [2,(3*5),4]`.
+   */
+  splitOp(sep: AluOp): IterableIterator<AluExp>;
   /**
    * Simplify the expression by replacing any known patterns and deduping
    * identical subexpressions.
@@ -208,10 +280,16 @@ declare class AluExp implements FpHashable {
   toString(): string;
   /** Generic fold() operation with a reducer over the expression tree. */
   fold<T = void>(reducer: (exp: AluExp, mappedSrc: T[]) => T): T;
+  /** Check if any expression in the tree satisfies a predicate. */
+  some(predicate: (exp: AluExp) => boolean): boolean;
   /** Rewrite the expression recursively using a visitor. */
   rewrite(visitor: (exp: AluExp) => AluExp | undefined | null): AluExp;
   /** Collect all nodes that satisfy a predicate. */
   collect(predicate: (exp: AluExp) => boolean): AluExp[];
+  /** Produce a list of all distinct AluOp in this expression. */
+  distinctOps(): Set<AluOp>;
+  /** Rewrite GlobalView operations to GlobalIndex operations. */
+  rewriteGlobalViews(): AluExp;
 }
 /** Symbolic form for each mathematical operation. */
 declare enum AluOp {
@@ -224,8 +302,11 @@ declare enum AluOp {
   Max = "Max",
   Sin = "Sin",
   Cos = "Cos",
+  Asin = "Asin",
+  Atan = "Atan",
   Exp = "Exp",
   Log = "Log",
+  Sqrt = "Sqrt",
   Reciprocal = "Reciprocal",
   Cast = "Cast",
   Bitcast = "Bitcast",
@@ -242,7 +323,7 @@ declare enum AluOp {
   Variable = "Variable",
   // arg = variable
   GlobalIndex = "GlobalIndex",
-  // arg = gid; src = [bufidx]
+  // arg = [gid, len]; src = [bufidx]
   GlobalView = "GlobalView",
 }
 /**
@@ -297,12 +378,12 @@ declare class Reduction implements FpHashable {
   /** Size of the reduction axis. */
   readonly size: number;
   /** Follow-up expression defined with the "acc" variable, defaults to identity. */
-  readonly fusion: AluExp;
+  readonly epilogue: AluExp;
   constructor(/** Data type of the values being reduced over. */
   dtype: DType, /** Operation to perform. Only ops in `AluGroup.Reduce` are supported. */
   op: AluOp, /** Size of the reduction axis. */
   size: number, /** Follow-up expression defined with the "acc" variable, defaults to identity. */
-  fusion?: AluExp);
+  epilogue?: AluExp);
   hash(state: FpHash): void;
   toString(): string;
   /** Get the identity for this reduction operation. */
@@ -313,10 +394,10 @@ declare class Reduction implements FpHashable {
 /** Expression for accessing `indices` in input array with the given shape. */
 //#endregion
 //#region src/backend.d.ts
-type Device = "cpu" | "webgpu";
+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.
  *
@@ -336,7 +417,7 @@ interface Backend {
   /** Maximum number of arguments per dispatched kernel. */
   readonly maxArgs: number;
   /** Allocate a new slot with reference count 1. */
-  malloc(size: number, initialData?: ArrayBuffer): Slot;
+  malloc(size: number, initialData?: Uint8Array): Slot;
   /** Increment the reference count of the slot. */
   incRef(slot: Slot): void;
   /**
@@ -345,9 +426,9 @@ interface Backend {
    */
   decRef(slot: Slot): void;
   /** Read a range of bytes from a buffer. */
-  read(slot: Slot, start?: number, count?: number): Promise<ArrayBuffer>;
+  read(slot: Slot, start?: number, count?: number): Promise<Uint8Array<ArrayBuffer>>;
   /** Read a range of bytes from a buffer, blocking variant. */
-  readSync(slot: Slot, start?: number, count?: number): ArrayBuffer;
+  readSync(slot: Slot, start?: number, count?: number): Uint8Array<ArrayBuffer>;
   /** Prepare an expression to be executed later. */
   prepare(kernel: Kernel): Promise<Executable>;
   /** Prepare an expression to be executed later, blocking variant. */
@@ -369,18 +450,22 @@ declare class Executable<T = any> {
   data: T);
 }
 declare namespace tree_d_exports {
-  export { JsTree, JsTreeDef, MapJsTree, NodeType, flatten, leaves, map, ref, structure, unflatten };
+  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 MapJsTree<T, A, B> = T extends A ? B : 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> };
-/** Analog to the JAX "pytree" object, but for JavaScript. */
+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 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;
@@ -406,6 +491,23 @@ declare function unflatten<T>(treedef: JsTreeDef, leaves: Iterable<T>): JsTree<T
 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 {
+  strides: number[];
+  padding: [number, number][];
+  lhsDilation: number[];
+  rhsDilation: number[];
+}
+/**
+ * Check that the shapes and parameters passed to convolution are valid.
+ *
+ * If the check succeeds, returns the output shape.
+ */
+
 //#endregion
 //#region src/frontend/core.d.ts
 /**
@@ -431,12 +533,20 @@ declare enum Primitive {
   RandomBits = "random_bits",
   Sin = "sin",
   Cos = "cos",
+  Asin = "asin",
+  Atan = "atan",
   Exp = "exp",
   Log = "log",
+  Sqrt = "sqrt",
   Min = "min",
   Max = "max",
   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",
   Transpose = "transpose",
@@ -448,7 +558,7 @@ declare enum Primitive {
   Gather = "gather",
   JitCall = "jit_call",
 }
-interface PrimitiveParamsImpl extends Record<Primitive, Record<string, unknown>> {
+interface PrimitiveParamsImpl extends Record<Primitive, Record<string, any>> {
   [Primitive.Cast]: {
     dtype: DType;
   };
@@ -459,6 +569,16 @@ interface PrimitiveParamsImpl extends Record<Primitive, Record<string, unknown>>
     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;
   };
@@ -480,10 +600,10 @@ interface PrimitiveParamsImpl extends Record<Primitive, Record<string, unknown>>
     axis: number[];
   };
   [Primitive.Shrink]: {
-    slice: [number, number][];
+    slice: Pair[];
   };
   [Primitive.Pad]: {
-    width: [number, number][];
+    width: Pair[];
   };
   [Primitive.Gather]: {
     axis: number[];
@@ -505,8 +625,10 @@ declare enum CompareOp {
   LessEqual = "less_equal",
 }
 /** @inline */
+type Axis = number | number[] | null;
+/** @inline */
 type ReduceOpts = {
-  keepDims?: boolean;
+  keepdims?: boolean;
 };
 type MainTrace = {
   level: number;
@@ -583,8 +705,13 @@ 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 the array. */
   get dtype(): DType;
+  /** The number of dimensions of the array. */
   get ndim(): number;
   /** @ignore */
   fullLower(): Tracer;
@@ -598,11 +725,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;
   /**
@@ -670,7 +797,7 @@ declare abstract class Tracer {
    * 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] | [number, number] | null | Tracer)[]): this;
+  slice(...index: (number | [] | [number] | Pair | null | Tracer)[]): this;
 }
 declare class ShapedArray implements AbstractValue {
   readonly shape: number[];
@@ -678,7 +805,7 @@ declare class ShapedArray implements AbstractValue {
   constructor(shape: number[], dtype: DType);
   static fromAval(aval: AbstractValue): ShapedArray;
   get ndim(): number;
-  strShort(): string;
+  toString(): string;
   equals(other: ShapedArray): boolean;
 }
 //#endregion
@@ -730,7 +857,11 @@ 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(source: AluExp | Slot, st: ShapeTracker, dtype: DType, backend: Backend, {
+    pending
+  }?: {
+    pending?: Iterable<PendingExecute> | null;
+  });
   /** @ignore */
   get aval(): ShapedArray;
   /** Return a simple string representation of the array's dimensions. */
@@ -749,14 +880,26 @@ declare class Array extends Tracer {
    */
   [Symbol.toPrimitive](): any;
   /** Realize the array and return it as data. */
-  data(): Promise<Float32Array | Int32Array | Uint32Array>;
-  /** Wait for this array to be placed on the backend, if needed. */
-  wait(): Promise<void>;
+  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(): Float32Array | Int32Array | Uint32Array;
+  dataSync(): DataArray;
   /**
    * Convert this array into a JavaScript object.
    *
@@ -769,17 +912,21 @@ declare class Array extends Tracer {
   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;
   _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 | Float32Array | Int32Array | RecursiveArray<number> | RecursiveArray<boolean>, {
+declare function array(values: Array | Float16Array<ArrayBuffer> | Float32Array<ArrayBuffer> | Int32Array<ArrayBuffer> | Uint32Array<ArrayBuffer> | RecursiveArray<number> | RecursiveArray<boolean>, {
   shape,
   dtype,
   device
@@ -815,7 +962,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
@@ -851,15 +998,14 @@ declare function linspace(start: number, stop: number, num?: number, endpoint?:
   dtype,
   device
 }?: DTypeAndDevice): Array;
-/** Translate a `CompareOp` into an `AluExp` on two sub-expressions. */
 declare namespace numpy_d_exports {
-  export { Array, ArrayLike, DType, abs, absolute, add, allclose, arange, argmax, argmin, array, astype, bool, clip, columnStack, complex64, concatenate, cos, diag, diagonal, divide, dot, dstack, e, equal, eulerGamma, exp, exp2, eye, flip, fliplr, flipud, float32, full, 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, pad, permuteDims, pi, prod, ravel, reciprocal, reshape, scalar, shape$1 as shape, sin, size, square, stack, sum, tan, transpose, trueDivide, trunc, uint32, vdot, vecdot, vstack, where, zeros };
+  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;
 declare const uint32 = DType.Uint32;
 declare const bool = DType.Bool;
-declare const complex64 = DType.Complex64;
+declare const float16 = DType.Float16;
 /** Euler's constant, `e = 2.7182818284590...` */
 declare const e: number;
 /** Euler-Mascheroni constant, `γ = 0.5772156649...` */
@@ -870,52 +1016,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;
-/** Return element-wise minimum of the input arrays. */
+/** @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;
-/** 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
@@ -923,10 +1083,28 @@ 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[];
+/**
+ * @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.
@@ -935,15 +1113,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.
  *
@@ -959,7 +1137,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.
  *
@@ -1003,16 +1181,45 @@ 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.
  *
  * 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.
+ * 3D or higher, compute diagonals along the two given axes (default: 0, 1).
  *
- * This returns a view over the existing array.
+ * 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;
 /**
@@ -1031,8 +1238,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.
  *
@@ -1050,6 +1277,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.
  *
@@ -1066,15 +1308,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;
@@ -1084,8 +1361,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 const sinh: OwnedFunction<(x: ArrayLike) => Array>;
+/**
+ * @function
+ * Calculate element-wise hyperbolic cosine of input.
+ *
+ * `cosh(x) = (exp(x) + exp(-x)) / 2`
+ */
+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 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;
@@ -1146,8 +1527,38 @@ declare class Jaxpr implements FpHashable {
   /** Flattens nested JitCall in a Jaxpr. Useful for handling jit-of-jit. */
   flatten(): Jaxpr;
 }
+/** @inline */
+type JitOpts = {
+  staticArgnums?: number[];
+  device?: Device;
+};
+declare namespace lax_d_exports {
+  export { PaddingType, conv, convGeneralDilated, convWithGeneralPadding, reduceWindow };
+}
+type PaddingType = "VALID" | "SAME" | "SAME_LOWER" | [number, number][];
+/**
+ * General n-dimensional convolution operator, with optional dilation.
+ *
+ * The semantics of this operation mimic the `jax.lax.conv_general_dilated`
+ * function in JAX, which wraps XLA's general convolution operator.
+ *
+ * Grouped convolutions are not supported right now.
+ */
+declare function convGeneralDilated(lhs: Array, rhs: Array, windowStrides: number[], padding: PaddingType, {
+  lhsDilation,
+  rhsDilation
+}?: {
+  lhsDilation?: number[];
+  rhsDilation?: 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;
+/** 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 { identity, logSigmoid, logSoftmax, logsumexp, 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:
@@ -1179,6 +1590,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))`.
@@ -1187,8 +1599,9 @@ declare function softSign(x: ArrayLike): Array;
  *
  * Reference: https://en.wikipedia.org/wiki/Swish_function
  */
-declare function 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))`.
@@ -1197,14 +1610,67 @@ declare function silu(x: ArrayLike): Array;
  *
  * Reference: https://en.wikipedia.org/wiki/Swish_function
  */
-declare const swish: typeof silu;
+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?: 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?: 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?: ArrayLike): Array;
+/**
+ * @function
+ * Gaussion error linear unit (GELU) activation function.
+ *
+ * This is computed element-wise. Currently jax-js does not support the erf() or
+ * gelu() functions exactly as primitives, so an approximation is used:
+ * `gelu(x) ~= x * 0.5 * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))`.
+ *
+ * Reference: https://ml-explore.github.io/mlx/build/html/python/nn/_autosummary_functions/mlx.nn.gelu_approx.html
+ *
+ * This will be improved in the future.
+ */
+declare const gelu: OwnedFunction<(x: ArrayLike) => Array>;
+/**
+ * Gated linear unit (GLU) activation function.
+ *
+ * Splits the `axis` dimension of the input into two halves, a and b, then
+ * 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.
+ *
+ * Computes the element-wise function:
+ * `mish(x) = x * tanh(softplus(x))`
+ */
+declare function mish(x: ArrayLike): Array;
 /**
  * Softmax function. Computes the function which rescales elements to the range
  * [0, 1] such that the elements along `axis` sum to 1.
@@ -1213,7 +1679,7 @@ declare const identity: (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.
  *
@@ -1222,7 +1688,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`.
  *
@@ -1231,7 +1697,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.
  *
@@ -1250,7 +1731,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;
@@ -1266,40 +1747,108 @@ declare function uniform(key: Array, shape?: number[], {
   minval?: number;
   maxval?: number;
 }): 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;
+/** Sample exponential random values according to `p(x) = exp(-x)`. */
+declare function exponential(key: Array, shape?: number[]): Array;
+/**
+ * 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 function normal(key: Array, shape?: number[]): Array;
 //#endregion
 //#region src/index.d.ts
-/** @inline */
-type WithArgsSubtype<F extends (args: any[]) => any, T> = Parameters<F> extends T ? F : never;
-/** Compute the forward-mode Jacobian-vector product for a function. */
-declare const jvp: <F extends (...args: any[]) => JsTree<Array>>(f: WithArgsSubtype<F, JsTree<ArrayLike>>, primals: MapJsTree<Parameters<F>, ArrayLike, ArrayLike>, tangents: MapJsTree<Parameters<F>, ArrayLike, ArrayLike>) => [ReturnType<F>, ReturnType<F>];
-/** Vectorize an operation on a batched axis for one or more inputs. */
-declare const vmap: <F extends (...args: any[]) => JsTree<Array>>(f: WithArgsSubtype<F, JsTree<ArrayLike>>, inAxes?: number | MapJsTree<Parameters<F>, ArrayLike, number | null>) => (...args: MapJsTree<Parameters<F>, ArrayLike, ArrayLike>) => ReturnType<F>;
-/** 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>, ArrayLike, ArrayLike>) => ReturnType<F>;
-/** Construct a Jaxpr by dynamically tracing a function with example inputs. */
-declare const makeJaxpr: <F extends (...args: any[]) => JsTree<Array>>(f: WithArgsSubtype<F, JsTree<ArrayLike>>) => (...args: Parameters<F>) => {
+/**
+ * @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>];
+/**
+ * @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>;
+/**
+ * @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>;
+/**
+ * @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;
 };
-declare const jit: <F extends (...args: any[]) => JsTree<Array>>(f: WithArgsSubtype<F, JsTree<ArrayLike>>) => (...args: MapJsTree<Parameters<F>, ArrayLike, ArrayLike>) => ReturnType<F>;
 /**
+ * @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,
+ *   and different values will trigger recompilation.
+ * - `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) => 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: WithArgsSubtype<F, JsTree<ArrayLike>>, ...primals: MapJsTree<Parameters<F>, ArrayLike, ArrayLike>) => [ReturnType<F>, (...tangents: MapJsTree<Parameters<F>, ArrayLike, ArrayLike>) => ReturnType<F>];
-/** Calculate the reverse-mode vector-Jacobian product for a function. */
-declare const vjp: <F extends (...args: any[]) => JsTree<Array>>(f: WithArgsSubtype<F, JsTree<ArrayLike>>, ...primals: MapJsTree<Parameters<F>, ArrayLike, ArrayLike>) => [ReturnType<F>, (cotangents: MapJsTree<ReturnType<F>, Array, ArrayLike>) => MapJsTree<Parameters<F>, ArrayLike, Array>];
+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>];
 /**
+ * @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: WithArgsSubtype<F, JsTree<ArrayLike>>) => (...primals: MapJsTree<Parameters<F>, ArrayLike, ArrayLike>) => MapJsTree<Parameters<F>[0], ArrayLike, Array>;
-/** Create a function that evaluates both `f` and the gradient of `f`. */
-declare const valueAndGrad: <F extends (...args: any[]) => JsTree<Array>>(f: WithArgsSubtype<F, JsTree<ArrayLike>>) => (...primals: MapJsTree<Parameters<F>, ArrayLike, ArrayLike>) => [ReturnType<F>, MapJsTree<Parameters<F>[0], ArrayLike, Array>];
-/** Compute the Jacobian evaluated row-by-row by reverse-mode AD. */
+declare const grad: <F extends (...args: any[]) => JsTree<Array>>(f: F) => (...primals: MapJsTree<Parameters<F>, Array, ArrayLike>) => MapJsTree<Parameters<F>[0], ArrayLike, Array>;
+/**
+ * @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>];
+/**
+ * @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()`. */
-declare const jacobian: <F extends (x: Array) => Array>(f: F) => (...args: MapJsTree<Parameters<F>, ArrayLike, ArrayLike>) => ReturnType<F>;
+/**
+ * @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 { type Device, devices, grad, init, jacfwd, jacobian, jacrev, jit, jvp, 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 };