@jax-js/jax 0.0.2 → 0.0.3

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,6 +109,17 @@ 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
@@ -130,13 +144,16 @@ declare class FpHash {
 /** 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>;
 /**
  * Mathematical expression on scalar values.
  *
@@ -162,6 +179,7 @@ declare class AluExp implements FpHashable {
   static cos(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 +190,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 +207,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 +240,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 {
@@ -226,6 +264,7 @@ declare enum AluOp {
   Cos = "Cos",
   Exp = "Exp",
   Log = "Log",
+  Sqrt = "Sqrt",
   Reciprocal = "Reciprocal",
   Cast = "Cast",
   Bitcast = "Bitcast",
@@ -242,7 +281,7 @@ declare enum AluOp {
   Variable = "Variable",
   // arg = variable
   GlobalIndex = "GlobalIndex",
-  // arg = gid; src = [bufidx]
+  // arg = [gid, len]; src = [bufidx]
   GlobalView = "GlobalView",
 }
 /**
@@ -297,12 +336,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,7 +352,7 @@ 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;
@@ -336,7 +375,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 +384,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 +408,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 +449,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
 /**
@@ -433,10 +493,16 @@ declare enum Primitive {
   Cos = "cos",
   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 +514,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 +525,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 +556,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[];
@@ -584,6 +660,7 @@ declare abstract class Tracer {
    */
   abstract dispose(): void;
   get shape(): number[];
+  get size(): number;
   get dtype(): DType;
   get ndim(): number;
   /** @ignore */
@@ -670,7 +747,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 +755,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
@@ -749,14 +826,23 @@ 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.
+   */
+  wait(): 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,6 +855,13 @@ 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;
@@ -779,7 +872,7 @@ declare function scalar(value: number | boolean, {
   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
@@ -851,15 +944,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, 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 };
 }
 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...` */
@@ -886,6 +978,8 @@ declare const cos: (x: ArrayLike) => Array;
 declare const exp: (x: ArrayLike) => Array;
 /** 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. */
+declare const sqrt: (x: ArrayLike) => Array;
 /** Return element-wise minimum of the input arrays. */
 declare const minimum: (x: ArrayLike, y: ArrayLike) => Array;
 /** Return element-wise maximum of the input arrays. */
@@ -927,6 +1021,12 @@ declare const pad: (x: ArrayLike, width: number | [number, number] | [number, nu
 declare const ndim: (x: ArrayLike) => number;
 /** 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. */
+declare const zerosLike: (a: ArrayLike, dtype?: DType) => Array;
+/** 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. */
+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.
@@ -1010,9 +1110,11 @@ declare function ravel(a: 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;
 /**
@@ -1084,6 +1186,24 @@ 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 element-wise hyperbolic sine of input.
+ *
+ * `sinh(x) = (exp(x) - exp(-x)) / 2`
+ */
+declare function sinh(x: ArrayLike): Array;
+/**
+ * Calculate element-wise hyperbolic cosine of input.
+ *
+ * `cosh(x) = (exp(x) + exp(-x)) / 2`
+ */
+declare function cosh(x: ArrayLike): Array;
+/**
+ * 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;
 //#endregion
 //#region src/frontend/jaxpr.d.ts
 /** Variable in a Jaxpr expression. */
@@ -1146,8 +1266,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, logsumexp, mish, oneHot, relu, relu6, sigmoid, silu, softSign, softmax, softplus, swish };
 }
 /**
  * Rectified Linear Unit (ReLU) activation function:
@@ -1187,7 +1337,7 @@ declare function softSign(x: ArrayLike): Array;
  *
  * Reference: https://en.wikipedia.org/wiki/Swish_function
  */
-declare function silu(x: ArrayLike): Array;
+declare const silu: (x: ArrayLike) => Array;
 /**
  * Sigmoid-weighted Linear Unit (SiLU) activation function, also known as
  * Swish, computed element-wise:
@@ -1197,7 +1347,7 @@ declare function silu(x: ArrayLike): Array;
  *
  * Reference: https://en.wikipedia.org/wiki/Swish_function
  */
-declare const swish: typeof silu;
+declare const swish: (x: ArrayLike) => Array;
 /**
  * Log-sigmoid activation function, computed element-wise:
  * `log_sigmoid(x) = log(sigmoid(x)) = -log(1 + exp(-x))`.
@@ -1205,6 +1355,48 @@ declare const swish: typeof silu;
 declare function logSigmoid(x: ArrayLike): Array;
 /** 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;
+/**
+ * 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;
+/**
+ * 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;
+/**
+ * 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: (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;
+/**
+ * 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.
@@ -1268,38 +1460,49 @@ declare function uniform(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>];
+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. */
-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>;
+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. */
-declare const jacfwd: <F extends (x: Array) => Array>(f: F) => (...args: MapJsTree<Parameters<F>, ArrayLike, ArrayLike>) => ReturnType<F>;
+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. */
-declare const makeJaxpr: <F extends (...args: any[]) => JsTree<Array>>(f: WithArgsSubtype<F, JsTree<ArrayLike>>) => (...args: Parameters<F>) => {
+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>;
+/**
+ * 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.
+ *
+ * **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) => (...args: MapJsTree<Parameters<F>, Array, ArrayLike>) => ReturnType<F>;
 /**
  * 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>];
+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. */
-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 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>];
 /**
  * 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>;
+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`. */
-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>];
+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. */
 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>;
+declare const jacobian: <F extends (x: Array) => Array>(f: F) => (...args: MapJsTree<Parameters<F>, Array, ArrayLike>) => ReturnType<F>;
 //#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 { 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 };