@jax-js/jax 0.0.4 → 0.1.0

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
 /**
@@ -177,12 +178,12 @@ type DataArray = Float32Array<ArrayBuffer> | Int32Array<ArrayBuffer> | Uint32Arr
  * **Type lattice:**
  * ```text
  * bool -> uint32 -> int32 -> float16 -> float32
- *  weak f* --^
+ *  weakType --^
  * ```
  *
- * The asterisk f* is a weak type used for JS number constants. When creating
- * arrays, JS numbers default to float32 but "weak" so they cast to the dtype of
- * any array they are first combined with.
+ * `weakType` represents weakly typed arrays. These are created for JS numbers,
+ * which default to float32 but "weak" so they cast to the dtype of any array
+ * they are first combined with, except `bool`.
  *
  * **Examples:**
  * - `promoteTypes(bool, int32) → int32`
@@ -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",
@@ -610,6 +619,7 @@ interface PrimitiveParamsImpl extends Record<Primitive, Record<string, any>> {
     outDim: number;
   };
   [Primitive.JitCall]: {
+    name: string;
     jaxpr: Jaxpr;
     numConsts: number;
   };
@@ -648,10 +658,40 @@ declare abstract class Trace {
   abstract lift(val: Tracer): Tracer;
   abstract processPrimitive<P extends Primitive>(primitive: P, tracers: Tracer[], params: PrimitiveParams<P>): Tracer[];
 }
+/** Internal representation of an array value. */
 interface AbstractValue {
+  /** Shape of the array. Must be a static tuple of non-negative dimensions. */
   shape: number[];
+  /** Concrete data type of array elements. */
   dtype: DType;
+  /**
+   * Arrays created from JavaScript numbers (e.g., `np.array(3)`) are created as
+   * _weakly typed_ unless a dtype is explicitly specified.
+   *
+   * Weakly typed values will automatically cast to the data type of other
+   * arrays when used as an operand as an expression. This property only affects
+   * how they promote in type casting; their memory layout is still determined
+   * by the actual `dtype` field.
+   *
+   * ```ts
+   * const x = np.array(3); // weakType = true, dtype = float32
+   * const y = np.array([1, 2], { dtype: np.int32 }); // weakType = false, dtype = int32
+   * const z = x.add(y); // z has dtype int32 because x is weakly typed
+   * ```
+   *
+   * Weak types are present in JIT programs in their spec (e.g., Jaxpr inputs
+   * and outputs can be weakly typed) form. But they're solely a frontend
+   * concept. Backends are not aware of weak types.
+   */
+  weakType: boolean;
 }
+/**
+ * Broadcast shapes and promote types with casting for two avals.
+ *
+ * This implements the weak type behavior described in `promoteTypes()`, but not
+ * implemented in that function as `weakType` is not passed.
+ */
+
 declare abstract class Tracer {
   /** @ignore */
   readonly _trace: Trace;
@@ -709,8 +749,15 @@ declare abstract class Tracer {
   get shape(): number[];
   /** The total number of elements in the array. */
   get size(): number;
-  /** The dtype of the array. */
+  /** The dtype of elements stored in the array. */
   get dtype(): DType;
+  /**
+   * Whether the array is weakly typed.
+   *
+   * Weakly typed arrays will cast to the dtype of the other operand. See
+   * `promoteTypes()` for details.
+   */
+  get weakType(): boolean;
   /** The number of dimensions of the array. */
   get ndim(): number;
   /** @ignore */
@@ -802,7 +849,8 @@ declare abstract class Tracer {
 declare class ShapedArray implements AbstractValue {
   readonly shape: number[];
   readonly dtype: DType;
-  constructor(shape: number[], dtype: DType);
+  readonly weakType: boolean;
+  constructor(shape: number[], dtype: DType, weakType: boolean);
   static fromAval(aval: AbstractValue): ShapedArray;
   get ndim(): number;
   toString(): string;
@@ -838,10 +886,19 @@ type DTypeAndDevice = {
   dtype?: DType;
   device?: Device;
 };
+type ArrayConstructorArgs = {
+  source: AluExp | Slot;
+  st: ShapeTracker;
+  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
@@ -857,11 +914,7 @@ declare class Array extends Tracer {
    * is a backend `Slot`, this constructor _takes ownership_ of the slot. It
    * will be freed when the array is disposed.
    */
-  constructor(source: AluExp | Slot, st: ShapeTracker, dtype: DType, backend: Backend, {
-    pending
-  }?: {
-    pending?: Iterable<PendingExecute> | null;
-  });
+  constructor(args: ArrayConstructorArgs);
   /** @ignore */
   get aval(): ShapedArray;
   /** Return a simple string representation of the array's dimensions. */
@@ -921,10 +974,13 @@ 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;
 }
-/** Construct an array from a single scalar constant. */
-
 /** Constructor for creating a new array from data. */
 declare function array(values: Array | Float16Array<ArrayBuffer> | Float32Array<ArrayBuffer> | Int32Array<ArrayBuffer> | Uint32Array<ArrayBuffer> | RecursiveArray<number> | RecursiveArray<boolean>, {
   shape,
@@ -999,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;
@@ -1082,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.
@@ -1312,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). */
@@ -1323,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>;
 /**
@@ -1456,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
 /**
@@ -1477,10 +1554,10 @@ declare class Var {
 }
 /** Literal in a Jaxpr expression. Currently, only scalars are supported. */
 declare class Lit {
-  readonly dtype: DType;
   readonly value: number;
   readonly aval: ShapedArray;
-  constructor(dtype: DType, value: number);
+  get dtype(): DType;
+  constructor(aval: AbstractValue, value: number);
 }
 type Atom = Var | Lit;
 declare class VarPrinter {
@@ -1530,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][];
 /**
@@ -1557,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 };
 }
@@ -1641,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.
  *
@@ -1730,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 };
 }
@@ -1739,14 +1835,14 @@ declare function key(seed: number): Array;
 declare function split(key: Array, num?: number | number[]): Array;
 /** Sample uniform bits in the form of unsigned integers. */
 declare function bits(key: Array, shape?: number[]): Array;
-/** Sample uniform random values in [minval, maxval) with given shape. */
-declare function uniform(key: Array, shape?: number[], {
-  minval,
-  maxval
-}?: {
-  minval?: number;
-  maxval?: number;
-}): Array;
+/**
+ * @function
+ * Sample uniform random values in [minval, maxval) with given shape.
+ */
+declare const uniform: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined, args_2?: {
+  minval?: number | undefined;
+  maxval?: number | undefined;
+} | undefined) => Array>;
 /**
  * Sample Bernoulli random variables with given mean (0,1 categorical).
  *
@@ -1754,16 +1850,29 @@ declare function uniform(key: Array, shape?: number[], {
  * and must be broadcastable to `shape`.
  */
 declare function bernoulli(key: Array, p?: ArrayLike, shape?: number[]): Array;
-/** Sample exponential random values according to `p(x) = exp(-x)`. */
-declare function exponential(key: Array, shape?: number[]): Array;
 /**
+ * @function
+ * Sample exponential random values according to `p(x) = exp(-x)`.
+ */
+declare const exponential: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined) => Array>;
+/**
+ * @function
  * Sample random values according to `p(x) = 1/sqrt(2pi) * exp(-x^2/2)`.
  *
  * Unlike JAX, this uses the Box-Muller transform. JAX uses the erf_inv primitive instead and
  * directly inverts the CDF, but we don't have support for that yet. Outputs will not be
  * bitwise identical to JAX.
  */
-declare function normal(key: Array, shape?: number[]): Array;
+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
 /**
@@ -1775,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.
@@ -1850,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