@jax-js/jax 0.0.5 → 0.1.0

package/dist/index.d.cts

@@ -124,6 +124,7 @@ declare class ShapeTracker {
   /** Like pad(), but allows for negative values. */
   padOrShrink(arg: Pair[]): ShapeTracker;
 }
+//# sourceMappingURL=shape.d.ts.map
 //#endregion
 //#region src/utils.d.ts
 /**
@@ -222,6 +223,8 @@ declare class AluExp implements FpHashable {
   static atan(a: AluExp): AluExp;
   static exp(a: AluExp): AluExp;
   static log(a: AluExp): AluExp;
+  static erf(a: AluExp): AluExp;
+  static erfc(a: AluExp): AluExp;
   static sqrt(a: AluExp): AluExp;
   static reciprocal(a: AluExp): AluExp;
   static cast(dtype: DType, a: AluExp): AluExp;
@@ -289,8 +292,8 @@ declare class AluExp implements FpHashable {
   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>;
+  /** Produce all distinct AluOp in this expression, with their dtypes. */
+  distinctOps(): Map<AluOp, Set<DType>>;
   /** Rewrite GlobalView operations to GlobalIndex operations. */
   rewriteGlobalViews(): AluExp;
 }
@@ -309,6 +312,8 @@ declare enum AluOp {
   Atan = "Atan",
   Exp = "Exp",
   Log = "Log",
+  Erf = "Erf",
+  Erfc = "Erfc",
   Sqrt = "Sqrt",
   Reciprocal = "Reciprocal",
   Cast = "Cast",
@@ -465,7 +470,7 @@ type JsTree<T> = T | JsTree<T>[] | {
   [key: string]: JsTree<T>;
 };
 type Same<X, Y> = (<T>() => T extends X ? 1 : 2) extends (<T>() => T extends Y ? 1 : 2) ? true : false;
-type MappedJsTree<T, A, B> = T extends A ? B : T extends 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> };
+type MappedJsTree<T, A, B> = T extends A ? B : T extends Array ? T : T extends globalThis.Array<infer U> ? number extends T["length"] ? MapJsTree<U, A, B>[] : { [K in keyof T]: MapJsTree<T[K], A, B> } : { [K in keyof T]: MapJsTree<T[K], A, B> };
 /** @ignore Convert a subtype of JsTree<A> into a JsTree<B>, with the same structure. */
 type MapJsTree<T, A, B> = Same<A, B> extends true ? T : MappedJsTree<T, A, B>;
 /** Represents the structure of a JsTree. */
@@ -477,6 +482,8 @@ declare class JsTreeDef {
   constructor(nodeType: NodeType, nodeMetadata: any,
   // Must be comparable with deepEqual.
   childTreedefs: JsTreeDef[]);
+  /** Get the total number of leaves in the tree. */
+  get size(): number;
   /** Returns a string representation of this tree definition. */
   toString(root?: boolean): string;
   /** Compare this tree definition with another. */
@@ -540,6 +547,8 @@ declare enum Primitive {
   Atan = "atan",
   Exp = "exp",
   Log = "log",
+  Erf = "erf",
+  Erfc = "erfc",
   Sqrt = "sqrt",
   Min = "min",
   Max = "max",
@@ -886,12 +895,13 @@ type ArrayConstructorArgs = {
   dtype: DType;
   weakType: boolean;
   backend: Backend;
+  committed: boolean;
   pending?: Iterable<PendingExecute>;
 };
 /**
  * A multidimensional numeric array with data stored on CPU or GPU.
  *
- * This is the library's core data type. Equivalent to `jnp.Array` from JAX, or
+ * This is the library's core data type. Equivalent to `jax.Array` from JAX, or
  * `torch.Tensor`.
  *
  * Not to be confused with the JavaScript "Array" constructor. Avoid importing
@@ -967,7 +977,12 @@ declare class Array extends Tracer {
   item(): number;
   /** @private Internal plumbing method for Array / Tracer ops. */
   static _implRules(): typeof implRules;
+  /** @private */
   _realizeSource(): number;
+  /** @private Put this array on a new backend, asynchronously. */
+  _put(backend: Backend): Promise<Array>;
+  /** @private Put this array on a new backend, synchronously. */
+  _putSync(backend: Backend): Array;
 }
 /** Constructor for creating a new array from data. */
 declare function array(values: Array | Float16Array<ArrayBuffer> | Float32Array<ArrayBuffer> | Int32Array<ArrayBuffer> | Uint32Array<ArrayBuffer> | RecursiveArray<number> | RecursiveArray<boolean>, {
@@ -1043,7 +1058,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, 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 };
+  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, hamming, hann, heaviside, 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;
@@ -1126,7 +1141,7 @@ declare const moveaxis: (x: ArrayLike, src: number, dst: number) => Array;
  * all axes are padded with the same width. Or if it is an array of pairs, each
  * pair specifies the padding for its corresponding axis.
  */
-declare const pad: (x: ArrayLike, width: number | [number, number] | [number, number][]) => Array;
+declare const pad: (x: ArrayLike, width: number | Pair | Pair[]) => Array;
 /**
  * @function
  * Return the number of dimensions of an array. Does not consume array reference.
@@ -1356,6 +1371,26 @@ declare function absolute(x: ArrayLike): Array;
 declare const abs: typeof absolute;
 /** Return an element-wise indication of sign of the input. */
 declare function sign(x: ArrayLike): Array;
+/**
+ * Return the Hamming window of size M, a taper with a weighted cosine bell.
+ *
+ * `w(n) = 0.54 - 0.46 * cos(2πn/(M-1))` for `0 <= n <= M-1`.
+ */
+declare function hamming(M: number): Array;
+/**
+ * Return the Hann window of size M, a taper with a weighted cosine bell.
+ *
+ * `w(n) = 0.5 - 0.5 * cos(2πn/(M-1))` for `0 <= n <= M-1`.
+ */
+declare function hann(M: number): Array;
+/**
+ * @function
+ * Compute the Heaviside step function. It is defined piecewise:
+ * - `heaviside(x1, x2) = 0` for `x1 < 0`,
+ * - `heaviside(x1, x2) = x2` for `x1 == 0`,
+ * - `heaviside(x1, x2) = 1` for `x1 > 0`.
+ */
+declare const heaviside: OwnedFunction<(x1: ArrayLike, x2: ArrayLike) => Array>;
 /** Calculate element-wise square of the input array. */
 declare function square(x: ArrayLike): Array;
 /** Element-wise tangent function (takes radians). */
@@ -1367,8 +1402,8 @@ declare function acos(x: ArrayLike): Array;
  * 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.
+ * stable than `sqrt(x1**2 + x2**2)`. We don't currently implement those
+ * stability improvements.
  */
 declare const hypot: OwnedFunction<(x1: ArrayLike, x2: ArrayLike) => Array>;
 /**
@@ -1500,6 +1535,7 @@ declare function std(x: ArrayLike, axis?: Axis, opts?: {
   mean?: ArrayLike;
   correction?: number;
 } & ReduceOpts): Array;
+//# sourceMappingURL=numpy.d.ts.map
 //#endregion
 //#region src/frontend/jaxpr.d.ts
 /**
@@ -1574,10 +1610,9 @@ declare class Jaxpr implements FpHashable {
 /** @inline */
 type JitOpts = {
   staticArgnums?: number[];
-  device?: Device;
 };
 declare namespace lax_d_exports {
-  export { PaddingType, conv, convGeneralDilated, convWithGeneralPadding, reduceWindow };
+  export { PaddingType, conv, convGeneralDilated, convWithGeneralPadding, erf, erfc, reduceWindow, stopGradient };
 }
 type PaddingType = "VALID" | "SAME" | "SAME_LOWER" | [number, number][];
 /**
@@ -1601,6 +1636,23 @@ declare function convWithGeneralPadding(lhs: Array, rhs: Array, windowStrides: n
 declare function conv(lhs: Array, rhs: Array, windowStrides: number[], padding: PaddingType): Array;
 /** Reduce a computation over padded windows. */
 declare function reduceWindow(operand: Array, computation: (x: Array) => Array, windowDimensions: number[], windowStrides?: number[]): Array;
+/** The error function: `erf(x) = 2/sqrt(pi) * int[0..x] exp(-t^2) dt`. */
+declare function erf(x: ArrayLike): Array;
+/**
+ * The complementary error function: `erfc(x) = 1 - erf(x)`.
+ *
+ * This function is more accurate than `1 - erf(x)` for large values of `x`,
+ * where `erf(x)` is very close to 1.
+ */
+declare function erfc(x: ArrayLike): Array;
+/**
+ * Stops gradient computation.
+ *
+ * Behaves as the identity function but prevents the flow of gradients during
+ * forward or reverse-mode automatic differentiation.
+ */
+declare function stopGradient(x: ArrayLike): Array;
+//# sourceMappingURL=lax.d.ts.map
 declare namespace nn_d_exports {
   export { celu, elu, gelu, glu, identity, leakyRelu, logSigmoid, logSoftmax, logmeanexp, logsumexp, mish, oneHot, relu, relu6, sigmoid, silu, softSign, softmax, softplus, squareplus, standardize, swish };
 }
@@ -1685,15 +1737,17 @@ 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)))`.
+ * This is computed element-wise. There are two variants depending on whether
+ * `approximate` is set (default true):
  *
- * Reference: https://ml-explore.github.io/mlx/build/html/python/nn/_autosummary_functions/mlx.nn.gelu_approx.html
+ * - Approximate: `gelu(x) ~= x * 0.5 * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))`
+ * - Exact: `gelu(x) = x * 0.5 * erfc(-x / sqrt(2))`
  *
- * This will be improved in the future.
+ * Reference: https://ml-explore.github.io/mlx/build/html/python/nn/_autosummary_functions/mlx.nn.gelu_approx.html
  */
-declare const gelu: OwnedFunction<(x: ArrayLike) => Array>;
+declare const gelu: OwnedFunction<(x: ArrayLike, opts?: {
+  approximate?: boolean | undefined;
+} | undefined) => Array>;
 /**
  * Gated linear unit (GLU) activation function.
  *
@@ -1774,6 +1828,7 @@ declare function standardize(x: ArrayLike, axis?: Axis, opts?: {
  * ```
  */
 declare function oneHot(x: Array, numClasses: number): Array;
+//# sourceMappingURL=nn.d.ts.map
 declare namespace random_d_exports {
   export { bernoulli, bits, exponential, key, normal, split, uniform };
 }
@@ -1812,6 +1867,15 @@ declare const exponential: OwnedFunction<(key: ArrayLike, shape?: number[] | und
  * bitwise identical to JAX.
  */
 declare const normal: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined) => Array>;
+//# sourceMappingURL=random.d.ts.map
+declare namespace scipy_special_d_exports {
+  export { erf, erfc, logSoftmax, logit, logsumexp, softmax };
+}
+/**
+ * @function
+ * The logit function, `logit(p) = log(p / (1-p))`.
+ */
+declare const logit: OwnedFunction<(x: ArrayLike) => Array>;
 //#endregion
 //#region src/index.d.ts
 /**
@@ -1823,7 +1887,7 @@ declare const jvp: <F extends (...args: any[]) => JsTree<Array>>(f: F, primals:
  * @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>;
+declare const vmap: <F extends (...args: any[]) => JsTree<Array>>(f: F, inAxes?: number | (number | null | JsTree<number | null>)[]) => (...args: MapJsTree<Parameters<F>, Array, ArrayLike>) => ReturnType<F>;
 /**
  * @function
  * Compute the Jacobian evaluated column-by-column by forward-mode AD.
@@ -1898,5 +1962,19 @@ declare const jacobian: <F extends (x: Array) => Array>(f: F) => (...args: MapJs
  * Does not consume reference to the arrays.
  */
 declare function blockUntilReady<T extends JsTree<any>>(x: T): Promise<T>;
+/**
+ * Transfer `x` to `device`.
+ *
+ * `x` may be a nested container of arrays or scalars. The resulting structure
+ * is committed to the device.
+ *
+ * If `device` is not specified, this function behaves as identity if the input
+ * is already an `Array`, otherwise it places the scalar uncommitted on the
+ * default device.
+ */
+declare function devicePut<T extends JsTree<any>>(x: T, device?: Device): Promise<MapJsTree<T, number | boolean, Array>>;
+//# sourceMappingURL=index.d.ts.map
+
 //#endregion
-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 };
+export { Array, DType, type Device, Jaxpr, type JsTree, type JsTreeDef, type OwnedFunction, blockUntilReady, defaultDevice, devicePut, 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, scipy_special_d_exports as scipySpecial, setDebug, tree_d_exports as tree, valueAndGrad, vjp, vmap };
+//# sourceMappingURL=index.d.cts.map

package/dist/index.d.ts

@@ -121,6 +121,7 @@ declare class ShapeTracker {
   /** Like pad(), but allows for negative values. */
   padOrShrink(arg: Pair[]): ShapeTracker;
 }
+//# sourceMappingURL=shape.d.ts.map
 //#endregion
 //#region src/utils.d.ts
 /**
@@ -219,6 +220,8 @@ declare class AluExp implements FpHashable {
   static atan(a: AluExp): AluExp;
   static exp(a: AluExp): AluExp;
   static log(a: AluExp): AluExp;
+  static erf(a: AluExp): AluExp;
+  static erfc(a: AluExp): AluExp;
   static sqrt(a: AluExp): AluExp;
   static reciprocal(a: AluExp): AluExp;
   static cast(dtype: DType, a: AluExp): AluExp;
@@ -286,8 +289,8 @@ declare class AluExp implements FpHashable {
   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>;
+  /** Produce all distinct AluOp in this expression, with their dtypes. */
+  distinctOps(): Map<AluOp, Set<DType>>;
   /** Rewrite GlobalView operations to GlobalIndex operations. */
   rewriteGlobalViews(): AluExp;
 }
@@ -306,6 +309,8 @@ declare enum AluOp {
   Atan = "Atan",
   Exp = "Exp",
   Log = "Log",
+  Erf = "Erf",
+  Erfc = "Erfc",
   Sqrt = "Sqrt",
   Reciprocal = "Reciprocal",
   Cast = "Cast",
@@ -462,7 +467,7 @@ type JsTree<T> = T | JsTree<T>[] | {
   [key: string]: JsTree<T>;
 };
 type Same<X, Y> = (<T>() => T extends X ? 1 : 2) extends (<T>() => T extends Y ? 1 : 2) ? true : false;
-type MappedJsTree<T, A, B> = T extends A ? B : T extends 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> };
+type MappedJsTree<T, A, B> = T extends A ? B : T extends Array ? T : T extends globalThis.Array<infer U> ? number extends T["length"] ? MapJsTree<U, A, B>[] : { [K in keyof T]: MapJsTree<T[K], A, B> } : { [K in keyof T]: MapJsTree<T[K], A, B> };
 /** @ignore Convert a subtype of JsTree<A> into a JsTree<B>, with the same structure. */
 type MapJsTree<T, A, B> = Same<A, B> extends true ? T : MappedJsTree<T, A, B>;
 /** Represents the structure of a JsTree. */
@@ -474,6 +479,8 @@ declare class JsTreeDef {
   constructor(nodeType: NodeType, nodeMetadata: any,
   // Must be comparable with deepEqual.
   childTreedefs: JsTreeDef[]);
+  /** Get the total number of leaves in the tree. */
+  get size(): number;
   /** Returns a string representation of this tree definition. */
   toString(root?: boolean): string;
   /** Compare this tree definition with another. */
@@ -537,6 +544,8 @@ declare enum Primitive {
   Atan = "atan",
   Exp = "exp",
   Log = "log",
+  Erf = "erf",
+  Erfc = "erfc",
   Sqrt = "sqrt",
   Min = "min",
   Max = "max",
@@ -883,12 +892,13 @@ type ArrayConstructorArgs = {
   dtype: DType;
   weakType: boolean;
   backend: Backend;
+  committed: boolean;
   pending?: Iterable<PendingExecute>;
 };
 /**
  * A multidimensional numeric array with data stored on CPU or GPU.
  *
- * This is the library's core data type. Equivalent to `jnp.Array` from JAX, or
+ * This is the library's core data type. Equivalent to `jax.Array` from JAX, or
  * `torch.Tensor`.
  *
  * Not to be confused with the JavaScript "Array" constructor. Avoid importing
@@ -964,7 +974,12 @@ declare class Array extends Tracer {
   item(): number;
   /** @private Internal plumbing method for Array / Tracer ops. */
   static _implRules(): typeof implRules;
+  /** @private */
   _realizeSource(): number;
+  /** @private Put this array on a new backend, asynchronously. */
+  _put(backend: Backend): Promise<Array>;
+  /** @private Put this array on a new backend, synchronously. */
+  _putSync(backend: Backend): Array;
 }
 /** Constructor for creating a new array from data. */
 declare function array(values: Array | Float16Array<ArrayBuffer> | Float32Array<ArrayBuffer> | Int32Array<ArrayBuffer> | Uint32Array<ArrayBuffer> | RecursiveArray<number> | RecursiveArray<boolean>, {
@@ -1040,7 +1055,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, 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 };
+  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, hamming, hann, heaviside, 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;
@@ -1123,7 +1138,7 @@ declare const moveaxis: (x: ArrayLike, src: number, dst: number) => Array;
  * all axes are padded with the same width. Or if it is an array of pairs, each
  * pair specifies the padding for its corresponding axis.
  */
-declare const pad: (x: ArrayLike, width: number | [number, number] | [number, number][]) => Array;
+declare const pad: (x: ArrayLike, width: number | Pair | Pair[]) => Array;
 /**
  * @function
  * Return the number of dimensions of an array. Does not consume array reference.
@@ -1353,6 +1368,26 @@ declare function absolute(x: ArrayLike): Array;
 declare const abs: typeof absolute;
 /** Return an element-wise indication of sign of the input. */
 declare function sign(x: ArrayLike): Array;
+/**
+ * Return the Hamming window of size M, a taper with a weighted cosine bell.
+ *
+ * `w(n) = 0.54 - 0.46 * cos(2πn/(M-1))` for `0 <= n <= M-1`.
+ */
+declare function hamming(M: number): Array;
+/**
+ * Return the Hann window of size M, a taper with a weighted cosine bell.
+ *
+ * `w(n) = 0.5 - 0.5 * cos(2πn/(M-1))` for `0 <= n <= M-1`.
+ */
+declare function hann(M: number): Array;
+/**
+ * @function
+ * Compute the Heaviside step function. It is defined piecewise:
+ * - `heaviside(x1, x2) = 0` for `x1 < 0`,
+ * - `heaviside(x1, x2) = x2` for `x1 == 0`,
+ * - `heaviside(x1, x2) = 1` for `x1 > 0`.
+ */
+declare const heaviside: OwnedFunction<(x1: ArrayLike, x2: ArrayLike) => Array>;
 /** Calculate element-wise square of the input array. */
 declare function square(x: ArrayLike): Array;
 /** Element-wise tangent function (takes radians). */
@@ -1364,8 +1399,8 @@ declare function acos(x: ArrayLike): Array;
  * 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.
+ * stable than `sqrt(x1**2 + x2**2)`. We don't currently implement those
+ * stability improvements.
  */
 declare const hypot: OwnedFunction<(x1: ArrayLike, x2: ArrayLike) => Array>;
 /**
@@ -1497,6 +1532,7 @@ declare function std(x: ArrayLike, axis?: Axis, opts?: {
   mean?: ArrayLike;
   correction?: number;
 } & ReduceOpts): Array;
+//# sourceMappingURL=numpy.d.ts.map
 //#endregion
 //#region src/frontend/jaxpr.d.ts
 /**
@@ -1571,10 +1607,9 @@ declare class Jaxpr implements FpHashable {
 /** @inline */
 type JitOpts = {
   staticArgnums?: number[];
-  device?: Device;
 };
 declare namespace lax_d_exports {
-  export { PaddingType, conv, convGeneralDilated, convWithGeneralPadding, reduceWindow };
+  export { PaddingType, conv, convGeneralDilated, convWithGeneralPadding, erf, erfc, reduceWindow, stopGradient };
 }
 type PaddingType = "VALID" | "SAME" | "SAME_LOWER" | [number, number][];
 /**
@@ -1598,6 +1633,23 @@ declare function convWithGeneralPadding(lhs: Array, rhs: Array, windowStrides: n
 declare function conv(lhs: Array, rhs: Array, windowStrides: number[], padding: PaddingType): Array;
 /** Reduce a computation over padded windows. */
 declare function reduceWindow(operand: Array, computation: (x: Array) => Array, windowDimensions: number[], windowStrides?: number[]): Array;
+/** The error function: `erf(x) = 2/sqrt(pi) * int[0..x] exp(-t^2) dt`. */
+declare function erf(x: ArrayLike): Array;
+/**
+ * The complementary error function: `erfc(x) = 1 - erf(x)`.
+ *
+ * This function is more accurate than `1 - erf(x)` for large values of `x`,
+ * where `erf(x)` is very close to 1.
+ */
+declare function erfc(x: ArrayLike): Array;
+/**
+ * Stops gradient computation.
+ *
+ * Behaves as the identity function but prevents the flow of gradients during
+ * forward or reverse-mode automatic differentiation.
+ */
+declare function stopGradient(x: ArrayLike): Array;
+//# sourceMappingURL=lax.d.ts.map
 declare namespace nn_d_exports {
   export { celu, elu, gelu, glu, identity, leakyRelu, logSigmoid, logSoftmax, logmeanexp, logsumexp, mish, oneHot, relu, relu6, sigmoid, silu, softSign, softmax, softplus, squareplus, standardize, swish };
 }
@@ -1682,15 +1734,17 @@ 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)))`.
+ * This is computed element-wise. There are two variants depending on whether
+ * `approximate` is set (default true):
  *
- * Reference: https://ml-explore.github.io/mlx/build/html/python/nn/_autosummary_functions/mlx.nn.gelu_approx.html
+ * - Approximate: `gelu(x) ~= x * 0.5 * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))`
+ * - Exact: `gelu(x) = x * 0.5 * erfc(-x / sqrt(2))`
  *
- * This will be improved in the future.
+ * Reference: https://ml-explore.github.io/mlx/build/html/python/nn/_autosummary_functions/mlx.nn.gelu_approx.html
  */
-declare const gelu: OwnedFunction<(x: ArrayLike) => Array>;
+declare const gelu: OwnedFunction<(x: ArrayLike, opts?: {
+  approximate?: boolean | undefined;
+} | undefined) => Array>;
 /**
  * Gated linear unit (GLU) activation function.
  *
@@ -1771,6 +1825,7 @@ declare function standardize(x: ArrayLike, axis?: Axis, opts?: {
  * ```
  */
 declare function oneHot(x: Array, numClasses: number): Array;
+//# sourceMappingURL=nn.d.ts.map
 declare namespace random_d_exports {
   export { bernoulli, bits, exponential, key, normal, split, uniform };
 }
@@ -1809,6 +1864,15 @@ declare const exponential: OwnedFunction<(key: ArrayLike, shape?: number[] | und
  * bitwise identical to JAX.
  */
 declare const normal: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined) => Array>;
+//# sourceMappingURL=random.d.ts.map
+declare namespace scipy_special_d_exports {
+  export { erf, erfc, logSoftmax, logit, logsumexp, softmax };
+}
+/**
+ * @function
+ * The logit function, `logit(p) = log(p / (1-p))`.
+ */
+declare const logit: OwnedFunction<(x: ArrayLike) => Array>;
 //#endregion
 //#region src/index.d.ts
 /**
@@ -1820,7 +1884,7 @@ declare const jvp: <F extends (...args: any[]) => JsTree<Array>>(f: F, primals:
  * @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>;
+declare const vmap: <F extends (...args: any[]) => JsTree<Array>>(f: F, inAxes?: number | (number | null | JsTree<number | null>)[]) => (...args: MapJsTree<Parameters<F>, Array, ArrayLike>) => ReturnType<F>;
 /**
  * @function
  * Compute the Jacobian evaluated column-by-column by forward-mode AD.
@@ -1895,5 +1959,19 @@ declare const jacobian: <F extends (x: Array) => Array>(f: F) => (...args: MapJs
  * Does not consume reference to the arrays.
  */
 declare function blockUntilReady<T extends JsTree<any>>(x: T): Promise<T>;
+/**
+ * Transfer `x` to `device`.
+ *
+ * `x` may be a nested container of arrays or scalars. The resulting structure
+ * is committed to the device.
+ *
+ * If `device` is not specified, this function behaves as identity if the input
+ * is already an `Array`, otherwise it places the scalar uncommitted on the
+ * default device.
+ */
+declare function devicePut<T extends JsTree<any>>(x: T, device?: Device): Promise<MapJsTree<T, number | boolean, Array>>;
+//# sourceMappingURL=index.d.ts.map
+
 //#endregion
-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 };
+export { Array, DType, type Device, Jaxpr, type JsTree, type JsTreeDef, type OwnedFunction, blockUntilReady, defaultDevice, devicePut, 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, scipy_special_d_exports as scipySpecial, setDebug, tree_d_exports as tree, valueAndGrad, vjp, vmap };
+//# sourceMappingURL=index.d.ts.map