@jax-js/jax 0.0.5 → 0.1.1

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
 /**
@@ -167,9 +168,10 @@ declare enum DType {
   Uint32 = "uint32",
   Bool = "bool",
   Float16 = "float16",
+  Float64 = "float64",
 }
 /** @inline */
-type DataArray = Float32Array<ArrayBuffer> | Int32Array<ArrayBuffer> | Uint32Array<ArrayBuffer> | Float16Array<ArrayBuffer>;
+type DataArray = Float32Array<ArrayBuffer> | Int32Array<ArrayBuffer> | Uint32Array<ArrayBuffer> | Float16Array<ArrayBuffer> | Float64Array<ArrayBuffer>;
 /**
  * Promote two dtypes to their join according to the type lattice.
  *
@@ -179,7 +181,7 @@ type DataArray = Float32Array<ArrayBuffer> | Int32Array<ArrayBuffer> | Uint32Arr
  *
  * **Type lattice:**
  * ```text
- * bool -> uint32 -> int32 -> float16 -> float32
+ * bool -> uint32 -> int32 -> float16 -> float32 -> float64
  *  weakType --^
  * ```
  *
@@ -222,6 +224,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;
@@ -240,6 +244,7 @@ declare class AluExp implements FpHashable {
   static u32(value: number): AluExp;
   static bool(value: boolean): AluExp;
   static f16(value: number): AluExp;
+  static f64(value: number): AluExp;
   not(): AluExp;
   /** Compute a reasonable expression hash with low collision rate. */
   getHash(): bigint;
@@ -289,8 +294,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 +314,8 @@ declare enum AluOp {
   Atan = "Atan",
   Exp = "Exp",
   Log = "Log",
+  Erf = "Erf",
+  Erfc = "Erfc",
   Sqrt = "Sqrt",
   Reciprocal = "Reciprocal",
   Cast = "Cast",
@@ -465,7 +472,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 +484,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 +549,8 @@ declare enum Primitive {
   Atan = "atan",
   Exp = "exp",
   Log = "log",
+  Erf = "erf",
+  Erfc = "erfc",
   Sqrt = "sqrt",
   Min = "min",
   Max = "max",
@@ -621,11 +632,9 @@ interface PrimitiveParamsImpl extends Record<Primitive, Record<string, any>> {
 /** Type of parameters taken by each primitive. */
 type PrimitiveParams<T extends Primitive> = T extends keyof PrimitiveParamsImpl ? PrimitiveParamsImpl[T] : Record<string, never>;
 declare enum CompareOp {
-  Greater = "greater",
   Less = "less",
   Equal = "equal",
   NotEqual = "not_equal",
-  GreaterEqual = "greater_equal",
   LessEqual = "less_equal",
 }
 /** @inline */
@@ -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,10 +977,15 @@ 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>, {
+declare function array(values: Array | DataArray | RecursiveArray<number> | RecursiveArray<boolean>, {
   shape,
   dtype,
   device
@@ -1043,13 +1058,14 @@ 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, float64, full, fullLike, greater, greaterEqual, hamming, hann, heaviside, hstack, hypot, identity$1 as identity, inf, inner, int32, isfinite, isinf, isnan, isneginf, isposinf, 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 float16 = DType.Float16;
+declare const float64 = DType.Float64;
 /** Euler's constant, `e = 2.7182818284590...` */
 declare const e: number;
 /** Euler-Mascheroni constant, `γ = 0.5772156649...` */
@@ -1126,7 +1142,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 +1372,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 +1403,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 +1536,20 @@ declare function std(x: ArrayLike, axis?: Axis, opts?: {
   mean?: ArrayLike;
   correction?: number;
 } & ReduceOpts): Array;
+/** Test element-wise for positive or negative infinity, return bool array. */
+declare function isinf(x: ArrayLike): Array;
+/** Test element-wise for NaN (Not a Number). */
+declare function isnan(x: ArrayLike): Array;
+/** Test element-wise for negative infinity, return bool array. */
+declare function isneginf(x: ArrayLike): Array;
+/** Test element-wise for positive infinity, return bool array. */
+declare function isposinf(x: ArrayLike): Array;
+/**
+ * @function
+ * Test element-wise for finite values (not infinity or NaN).
+ */
+declare const isfinite: OwnedFunction<(x: ArrayLike) => Array>;
+//# sourceMappingURL=numpy.d.ts.map
 //#endregion
 //#region src/frontend/jaxpr.d.ts
 /**
@@ -1574,10 +1624,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 +1650,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 +1751,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 +1842,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 +1881,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 +1901,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 +1976,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
 /**
@@ -164,9 +165,10 @@ declare enum DType {
   Uint32 = "uint32",
   Bool = "bool",
   Float16 = "float16",
+  Float64 = "float64",
 }
 /** @inline */
-type DataArray = Float32Array<ArrayBuffer> | Int32Array<ArrayBuffer> | Uint32Array<ArrayBuffer> | Float16Array<ArrayBuffer>;
+type DataArray = Float32Array<ArrayBuffer> | Int32Array<ArrayBuffer> | Uint32Array<ArrayBuffer> | Float16Array<ArrayBuffer> | Float64Array<ArrayBuffer>;
 /**
  * Promote two dtypes to their join according to the type lattice.
  *
@@ -176,7 +178,7 @@ type DataArray = Float32Array<ArrayBuffer> | Int32Array<ArrayBuffer> | Uint32Arr
  *
  * **Type lattice:**
  * ```text
- * bool -> uint32 -> int32 -> float16 -> float32
+ * bool -> uint32 -> int32 -> float16 -> float32 -> float64
  *  weakType --^
  * ```
  *
@@ -219,6 +221,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;
@@ -237,6 +241,7 @@ declare class AluExp implements FpHashable {
   static u32(value: number): AluExp;
   static bool(value: boolean): AluExp;
   static f16(value: number): AluExp;
+  static f64(value: number): AluExp;
   not(): AluExp;
   /** Compute a reasonable expression hash with low collision rate. */
   getHash(): bigint;
@@ -286,8 +291,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 +311,8 @@ declare enum AluOp {
   Atan = "Atan",
   Exp = "Exp",
   Log = "Log",
+  Erf = "Erf",
+  Erfc = "Erfc",
   Sqrt = "Sqrt",
   Reciprocal = "Reciprocal",
   Cast = "Cast",
@@ -462,7 +469,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 +481,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 +546,8 @@ declare enum Primitive {
   Atan = "atan",
   Exp = "exp",
   Log = "log",
+  Erf = "erf",
+  Erfc = "erfc",
   Sqrt = "sqrt",
   Min = "min",
   Max = "max",
@@ -618,11 +629,9 @@ interface PrimitiveParamsImpl extends Record<Primitive, Record<string, any>> {
 /** Type of parameters taken by each primitive. */
 type PrimitiveParams<T extends Primitive> = T extends keyof PrimitiveParamsImpl ? PrimitiveParamsImpl[T] : Record<string, never>;
 declare enum CompareOp {
-  Greater = "greater",
   Less = "less",
   Equal = "equal",
   NotEqual = "not_equal",
-  GreaterEqual = "greater_equal",
   LessEqual = "less_equal",
 }
 /** @inline */
@@ -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,10 +974,15 @@ 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>, {
+declare function array(values: Array | DataArray | RecursiveArray<number> | RecursiveArray<boolean>, {
   shape,
   dtype,
   device
@@ -1040,13 +1055,14 @@ 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, float64, full, fullLike, greater, greaterEqual, hamming, hann, heaviside, hstack, hypot, identity$1 as identity, inf, inner, int32, isfinite, isinf, isnan, isneginf, isposinf, 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 float16 = DType.Float16;
+declare const float64 = DType.Float64;
 /** Euler's constant, `e = 2.7182818284590...` */
 declare const e: number;
 /** Euler-Mascheroni constant, `γ = 0.5772156649...` */
@@ -1123,7 +1139,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 +1369,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 +1400,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 +1533,20 @@ declare function std(x: ArrayLike, axis?: Axis, opts?: {
   mean?: ArrayLike;
   correction?: number;
 } & ReduceOpts): Array;
+/** Test element-wise for positive or negative infinity, return bool array. */
+declare function isinf(x: ArrayLike): Array;
+/** Test element-wise for NaN (Not a Number). */
+declare function isnan(x: ArrayLike): Array;
+/** Test element-wise for negative infinity, return bool array. */
+declare function isneginf(x: ArrayLike): Array;
+/** Test element-wise for positive infinity, return bool array. */
+declare function isposinf(x: ArrayLike): Array;
+/**
+ * @function
+ * Test element-wise for finite values (not infinity or NaN).
+ */
+declare const isfinite: OwnedFunction<(x: ArrayLike) => Array>;
+//# sourceMappingURL=numpy.d.ts.map
 //#endregion
 //#region src/frontend/jaxpr.d.ts
 /**
@@ -1571,10 +1621,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 +1647,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 +1748,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 +1839,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 +1878,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 +1898,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 +1973,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