@jax-js/jax 0.1.8 → 0.1.10

package/dist/index.d.ts

@@ -433,9 +433,14 @@ declare class Routine {
 }
 /** One of the valid `Routine` that can be dispatched to backend. */
 declare enum Routines {
-  /** Stable sorting algorithm along the last axis. */
+  /**
+   * Sort along the last axis.
+   *
+   * This may be _unstable_ but it often doesn't matter, sorting numbers is
+   * bitwise unique up to signed zeros and NaNs.
+   */
   Sort = "Sort",
-  /** Returns `int32` indices of the stably sorted array. */
+  /** Stable sorting, returns `int32` indices and values of the sorted array. */
   Argsort = "Argsort",
   /**
    * Solve a triangular system of equations.
@@ -747,9 +752,9 @@ declare enum Primitive {
   Shrink = "shrink",
   Pad = "pad",
   Sort = "sort",
-  // sort(x, axis=-1)
+  // sort(x, axis=-1), unstable
   Argsort = "argsort",
-  // argsort(x, axis=-1)
+  // argsort(x, axis=-1), stable
   TriangularSolve = "triangular_solve",
   // A is upper triangular, A @ X.T = B.T
   Cholesky = "cholesky",
@@ -996,6 +1001,8 @@ declare abstract class Tracer {
   reshape(shape: number | number[]): this;
   /** Copy the array and cast to a specified dtype. */
   astype(dtype: DType): this;
+  /** Return a bitwise cast of the array, viewed as a new dtype. */
+  view(dtype?: DType): this;
   /** Subtract an array from this one. */
   sub(other: this | TracerValue): this;
   /** Divide an array by this one. */
@@ -1026,8 +1033,9 @@ declare abstract class Tracer {
    */
   sort(axis?: number): this;
   /**
-   * Return the indices that would sort an array. This may not be a stable
-   * sorting algorithm; it need not preserve order of indices in ties.
+   * Return the indices that would sort an array. Unlike `sort`, this is
+   * guaranteed to be a stable sorting algorithm; it always returns the smaller
+   * index first in event of ties.
    *
    * See `jax.numpy.argsort` for full docs.
    */
@@ -1109,6 +1117,12 @@ type DTypeAndDevice = {
   dtype?: DType;
   device?: Device;
 };
+/** @inline */
+type DTypeShapeAndDevice = {
+  dtype?: DType;
+  shape?: number[];
+  device?: Device;
+};
 type ArrayConstructorArgs = {
   source: AluExp | Slot;
   st: ShapeTracker;
@@ -1218,15 +1232,9 @@ declare function array(values: Array | DataArray | RecursiveArray<number> | Recu
 type ImplRule<P extends Primitive> = (tracers: Array[], params: PrimitiveParams<P>) => Array[];
 declare const implRules: { [P in Primitive]: ImplRule<P> };
 /** Return a new array of given shape and type, filled with zeros. */
-declare function zeros(shape: number[], {
-  dtype,
-  device
-}?: DTypeAndDevice): Array;
+declare function zeros(shape: number[], opts?: DTypeAndDevice): Array;
 /** Return a new array of given shape and type, filled with ones. */
-declare function ones(shape: number[], {
-  dtype,
-  device
-}?: DTypeAndDevice): Array;
+declare function ones(shape: number[], opts?: DTypeAndDevice): Array;
 /** Return a new array of given shape and type, filled with `fill_value`. */
 declare function full(shape: number[], fillValue: number | boolean | Array, {
   dtype,
@@ -1418,8 +1426,10 @@ declare function triangularSolve(a: ArrayLike, b: ArrayLike, {
   unitDiagonal?: boolean;
 }): Array;
 declare namespace lax_d_exports {
-  export { DotDimensionNumbers, PaddingType, conv, convGeneralDilated, convTranspose, convWithGeneralPadding, dot$1 as dot, erf, erfc, lax_linalg_d_exports as linalg, reduceWindow, stopGradient };
+  export { DotDimensionNumbers, PaddingType, bitcastConvertType, conv, convGeneralDilated, convTranspose, convWithGeneralPadding, dot$1 as dot, erf, erfc, lax_linalg_d_exports as linalg, reduceWindow, stopGradient, topK };
 }
+/** Elementwise bitcast an array into a new dtype. */
+declare function bitcastConvertType(x: ArrayLike, newDtype: DType): Array;
 /**
  * Dimension numbers for general `dot()` primitive.
  *
@@ -1524,6 +1534,16 @@ declare function erfc(x: ArrayLike): Array;
  * forward or reverse-mode automatic differentiation.
  */
 declare function stopGradient(x: ArrayLike): Array;
+/**
+ * Returns top `k` values and their indices along the specified axis of operand.
+ *
+ * This is a _stable_ algorithm: If two elements are equal, the lower-index
+ * element appears first.
+ *
+ * @returns A tuple of `(values, indices)`, where `values` and `indices` have
+ * the same shape as `x`, except along `axis` where they have size `k`.
+ */
+declare function topK(x: ArrayLike, k: number, axis?: number): [Array, Array];
 declare namespace numpy_fft_d_exports {
   export { ComplexPair, fft, ifft };
 }
@@ -1548,7 +1568,7 @@ declare function fft(a: ComplexPair, axis?: number): ComplexPair;
  */
 declare function ifft(a: ComplexPair, axis?: number): ComplexPair;
 declare namespace numpy_linalg_d_exports {
-  export { cholesky, det, diagonal, inv, lstsq, matmul, matrixPower, matrixTranspose, outer, slogdet, solve, tensordot, trace, vecdot };
+  export { cholesky, cross$1 as cross, det, diagonal, inv, lstsq, matmul, matrixPower, matrixTranspose, outer, slogdet, solve, tensordot, trace, vecdot };
 }
 /**
  * Compute the Cholesky decomposition of a (batched) positive-definite matrix.
@@ -1563,6 +1583,13 @@ declare function cholesky(a: ArrayLike, {
   upper?: boolean;
   symmetrizeInput?: boolean;
 }): Array;
+/**
+ * Compute the cross-product of two 3D vectors.
+ *
+ * This is a simpler and less flexible version of `jax.numpy.cross()`.
+ * Both inputs must have size 3 along the specified axis.
+ */
+declare function cross$1(x1: ArrayLike, x2: ArrayLike, axis?: number): Array;
 /** Compute the determinant of a square matrix (batched). */
 declare function det(a: ArrayLike): Array;
 /** Compute the inverse of a square matrix (batched). */
@@ -1649,7 +1676,7 @@ type IInfo = Readonly<{
 /** Machine limits for integer types. */
 declare function iinfo(dtype: DType): IInfo;
 declare namespace numpy_d_exports {
-  export { Array, ArrayLike, DType, absolute as abs, absolute, acos, arccosh as acosh, add, all, allclose, any, arange, acos as arccos, arccosh, asin as arcsin, arcsinh, atan as arctan, atan2 as arctan2, arctanh, argmax, argmin, argsort, array, asin, arcsinh as asinh, astype, atan, atan2, arctanh as atanh, bool, broadcastArrays, broadcastShapes, broadcastTo, cbrt, ceil, clip, columnStack, concatenate, convolve, corrcoef, correlate, cos, cosh, cov, cumsum, cumsum as cumulativeSum, deg2rad, degrees, diag, diagonal, trueDivide as divide, divmod, dot, dstack, e, einsum, equal, eulerGamma, exp, exp2, expandDims, expm1, eye, numpy_fft_d_exports as fft, finfo, flip, fliplr, flipud, float16, float32, float64, floor, floorDivide, fmod, frexp, full, fullLike, greater, greaterEqual, hamming, hann, heaviside, hstack, hypot, identity$1 as identity, iinfo, inf, inner, int32, isfinite, isinf, isnan, isneginf, isposinf, ldexp, less, lessEqual, numpy_linalg_d_exports as linalg, linspace, log, log10, log1p, log2, logspace, matmul, matrixTranspose, max, maximum, mean, meshgrid, min, minimum, moveaxis, multiply, nan, nanToNum, ndim, negative, notEqual, ones, onesLike, outer, pad, transpose as permuteDims, pi, positive, power as pow, power, prod, promoteTypes, ptp, rad2deg, radians, ravel, reciprocal, remainder, repeat, reshape, shape$1 as shape, sign, sin, sinc, sinh, size, sort, split$1 as split, sqrt, square, squeeze, stack, std, subtract, sum, swapaxes, take, tan, tanh, tensordot, tile, trace, transpose, tri, tril, triu, trueDivide, trunc, uint32, var_, vdot, vecdot, vstack, where, zeros, zerosLike };
+  export { Array, ArrayLike, DType, absolute as abs, absolute, acos, arccosh as acosh, add, all, allclose, any, arange, acos as arccos, arccosh, asin as arcsin, arcsinh, atan as arctan, atan2 as arctan2, arctanh, argmax, argmin, argsort, round as around, array, arrayEqual, arrayEquiv, asin, arcsinh as asinh, astype, atan, atan2, arctanh as atanh, average, bool, broadcastArrays, broadcastShapes, broadcastTo, cbrt, ceil, clip, columnStack, concatenate, convolve, copysign, corrcoef, correlate, cos, cosh, cov, cross, cumsum, cumsum as cumulativeSum, deg2rad, degrees, diag, diagonal, trueDivide as divide, divmod, dot, dstack, e, einsum, equal, eulerGamma, exp, exp2, expandDims, expm1, eye, numpy_fft_d_exports as fft, finfo, flip, fliplr, flipud, float16, float32, float64, floor, floorDivide, fmod, frexp, full, fullLike, greater, greaterEqual, hann, heaviside, hstack, hypot, identity$1 as identity, iinfo, inf, inner, int32, isfinite, isinf, isnan, isneginf, isposinf, ldexp, less, lessEqual, numpy_linalg_d_exports as linalg, linspace, log, log10, log1p, log2, logicalAnd, logicalNot, logicalOr, logicalXor, logspace, matmul, matrixTranspose, matvec, max, maximum, mean, meshgrid, min, minimum, moveaxis, multiply, nan, nanToNum, ndim, negative, notEqual, ones, onesLike, outer, pad, transpose as permuteDims, pi, positive, power as pow, power, prod, promoteTypes, ptp, rad2deg, radians, ravel, reciprocal, remainder, repeat, reshape, rint, round, shape$1 as shape, sign, sin, sinc, sinh, size, sort, split$1 as split, sqrt, square, squeeze, stack, std, subtract, sum, swapaxes, take, tan, tanh, tensordot, tile, trace, transpose, tri, tril, triu, trueDivide, trunc, uint32, var_, vdot, vecdot, vecmat, vstack, where, zeros, zerosLike };
 }
 declare const float32 = DType.Float32;
 declare const int32 = DType.Int32;
@@ -1709,6 +1736,14 @@ declare const notEqual: (x: ArrayLike, y: ArrayLike) => Array;
 declare const greaterEqual: (x: ArrayLike, y: ArrayLike) => Array;
 /** @function Compare two arrays element-wise. */
 declare const lessEqual: (x: ArrayLike, y: ArrayLike) => Array;
+/** Compute element-wise logical AND. */
+declare function logicalAnd(x: ArrayLike, y: ArrayLike): Array;
+/** Compute element-wise logical OR. */
+declare function logicalOr(x: ArrayLike, y: ArrayLike): Array;
+/** Compute element-wise logical XOR. */
+declare function logicalXor(x: ArrayLike, y: ArrayLike): Array;
+/** Compute element-wise logical NOT. */
+declare function logicalNot(x: ArrayLike): Array;
 /** @function Element-wise ternary operator, evaluates to `x` if cond else `y`. */
 declare const where: (cond: ArrayLike, x: ArrayLike, y: ArrayLike) => Array;
 /**
@@ -1749,17 +1784,17 @@ declare const shape$1: (x: ArrayLike) => number[];
  * @function
  * Return an array of zeros with the same shape and type as a given array.
  */
-declare const zerosLike: (a: ArrayLike, dtype?: DType) => Array;
+declare const zerosLike: (a: ArrayLike, opts?: DTypeShapeAndDevice) => Array;
 /**
  * @function
  * Return an array of ones with the same shape and type as a given array.
  */
-declare const onesLike: (a: ArrayLike, dtype?: DType) => Array;
+declare const onesLike: (a: ArrayLike, opts?: DTypeShapeAndDevice) => Array;
 /**
  * @function
  * Return a full array with the same shape and type as a given array.
  */
-declare const fullLike: (a: ArrayLike, fillValue: number | boolean | Array, dtype?: DType) => Array;
+declare const fullLike: (a: ArrayLike, fillValue: number | boolean | Array, opts?: DTypeShapeAndDevice) => Array;
 /**
  * Return the number of elements in an array, optionally along an axis.
  * Does not consume array reference.
@@ -1793,6 +1828,16 @@ declare function all(a: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
 declare function ptp(a: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
 /** Compute the average of the array elements along the specified axis. */
 declare function mean(a: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
+/**
+ * Compute the weighted average along the specified axis.
+ *
+ * If no axis is specified, mean is computed along all the axes. The weights
+ * should have shape matching that of `a`, or if an axis is specified, it should
+ * match the shape along those axes.
+ */
+declare function average(a: ArrayLike, axis?: Axis, opts?: {
+  weights?: ArrayLike;
+} & ReduceOpts): Array;
 /**
  * Returns the indices of the minimum values along an axis.
  *
@@ -1948,8 +1993,9 @@ declare function trace(a: ArrayLike, offset?: number, axis1?: number, axis2?: nu
  */
 declare function sort(a: ArrayLike, axis?: number): Array;
 /**
- * Return indices that would sort an array. This may be an unstable sorting
- * algorithm; it need not preserve order of indices in ties.
+ * Return indices that would sort an array. Unlike `sort`, this is guaranteed to
+ * be a stable sorting algorithm; it always returns the smaller index first in
+ * event of ties.
  *
  * Returns an array of `int32` indices.
  *
@@ -1963,13 +2009,39 @@ declare function argsort(a: ArrayLike, axis?: number): Array;
  * numbered axis. By default, the flattened array is used.
  */
 declare function take(a: ArrayLike, indices: ArrayLike, axis?: number | null): Array;
-/** Return if two arrays are element-wise equal within a tolerance. */
+/**
+ * Return if two arrays are element-wise equal within a tolerance.
+ *
+ * The formula used is `|actual - expected| <= atol + rtol * |expected|`, with
+ * NaN values comparing equal if `equalNaN` is true.
+ */
 declare function allclose(actual: Parameters<typeof array>[0], expected: Parameters<typeof array>[0], options?: {
   rtol?: number;
   atol?: number;
+  equalNaN?: boolean;
 }): boolean;
+/**
+ * Check if two arrays are element-wise equal.
+ *
+ * Returns False if the arrays have different shapes. If `equalNaN` is True,
+ * NaNs in the same position are considered equal.
+ */
+declare function arrayEqual(a1: ArrayLike, a2: ArrayLike, opts?: {
+  equalNaN?: boolean;
+}): Array;
+/**
+ * Check if two arrays are element-wise equal after broadcasting.
+ *
+ * Unlike `arrayEqual`, this allows inputs with different but
+ * broadcast-compatible shapes.
+ */
+declare function arrayEquiv(a1: ArrayLike, a2: ArrayLike): Array;
 /** Matrix product of two arrays. */
 declare function matmul(x: ArrayLike, y: ArrayLike): Array;
+/** Matrix-vector product. x1 is [..., M, N], x2 is [..., N] → [..., M]. */
+declare function matvec(x1: ArrayLike, x2: ArrayLike): Array;
+/** Vector-matrix product. x1 is [..., N], x2 is [..., N, M] → [..., M]. */
+declare function vecmat(x1: ArrayLike, x2: ArrayLike): Array;
 /** Dot product of two arrays. */
 declare function dot(x: ArrayLike, y: ArrayLike): Array;
 /**
@@ -2022,6 +2094,18 @@ declare function inner(x: ArrayLike, y: ArrayLike): Array;
  * be of shape `[x.size, y.size]`.
  */
 declare function outer(x: ArrayLike, y: ArrayLike): Array;
+/**
+ * @function Compute the cross product of two arrays.
+ *
+ * Supports 2D (scalar result) and 3D cross products, with optional axis
+ * arguments. If `axis` is given, it overrides `axisa`, `axisb`, and `axisc`.
+ */
+declare const cross: OwnedFunction<(a: ArrayLike, b: ArrayLike, args_2?: {
+  axisa?: number | undefined;
+  axisb?: number | undefined;
+  axisc?: number | undefined;
+  axis?: number | undefined;
+} | undefined) => Array>;
 /** Vector dot product of two arrays along a given axis. */
 declare function vecdot(x: ArrayLike, y: ArrayLike, {
   axis
@@ -2067,14 +2151,13 @@ declare function clip(a: ArrayLike, min?: ArrayLike, max?: ArrayLike): Array;
 declare function absolute(x: ArrayLike): Array;
 /** Return an element-wise indication of sign of the input. */
 declare function sign(x: ArrayLike): Array;
-/** @function Return element-wise positive values of the input (no-op). */
-declare const positive: (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`.
+ * @function
+ * Return the value with the magnitude of x and the sign of y, element-wise.
  */
-declare function hamming(M: number): Array;
+declare const copysign: OwnedFunction<(x: ArrayLike, y: ArrayLike) => Array>;
+/** @function Return element-wise positive values of the input (no-op). */
+declare const positive: (x: ArrayLike) => Array;
 /**
  * Return the Hann window of size M, a taper with a weighted cosine bell.
  *
@@ -2169,6 +2252,18 @@ declare const remainder: OwnedFunction<(x: ArrayLike, y: ArrayLike) => Array>;
 declare function divmod(x: ArrayLike, y: ArrayLike): [Array, Array];
 /** Round input to the nearest integer towards zero. */
 declare function trunc(x: ArrayLike): Array;
+/**
+ * @function
+ * Round to the given number of decimals.
+ *
+ * Uses banker's rounding (round half to even) to match NumPy/JAX behavior.
+ */
+declare const round: OwnedFunction<(a: ArrayLike, decimals?: number | undefined) => Array>;
+/**
+ * @function
+ * Round to the nearest integer, with ties going to the nearest even integer.
+ */
+declare const rint: OwnedFunction<(x: ArrayLike) => Array>;
 /**
  * Compute `x1 * 2 ** x2` as a standard multiplication and exponentiation.
  *
@@ -2561,7 +2656,7 @@ declare function dotProductAttention(query: ArrayLike, key: ArrayLike, value: Ar
   localWindowSize?: number | [number, number];
 }): Array;
 declare namespace random_d_exports {
-  export { bernoulli, bits, cauchy, exponential, gumbel, key, laplace, multivariateNormal, normal, split, uniform };
+  export { bernoulli, bits, categorical, cauchy, exponential, gumbel, key, laplace, multivariateNormal, normal, split, uniform };
 }
 /** Create a pseudo-random number generator (PRNG) key from 32-bit integer seed. */
 declare function key(seed: ArrayLike): Array;
@@ -2584,6 +2679,32 @@ declare const uniform: OwnedFunction<(key: ArrayLike, shape?: number[] | undefin
  * and must be broadcastable to `shape`.
  */
 declare function bernoulli(key: Array, p?: ArrayLike, shape?: number[]): Array;
+/**
+ * @function
+ * Sample random values from categorical distributions.
+ *
+ * Uses the Gumbel max trick for sampling with replacement, or the Gumbel top-k
+ * trick for sampling without replacement.
+ *
+ * Note: Sampling without replacement currently uses argsort and slices the last
+ * k elements. This should be replaced with a more efficient topK implementation.
+ *
+ * - `key` - PRNG key
+ * - `logits` - Unnormalized log probabilities of the categorical distribution(s).
+ *   `softmax(logits, axis)` gives the corresponding probabilities.
+ * - `axis` - Axis along which logits belong to the same categorical distribution.
+ * - `shape` - Result batch shape. Must be broadcast-compatible with
+ *   `logits.shape` with `axis` removed. Default is `logits.shape` with `axis` removed.
+ * - `replace` - If true (default), sample with replacement. If false, sample
+ *   without replacement (each category can only be selected once per batch).
+ * @returns A random array with int dtype and shape given by `shape` if provided,
+ *   otherwise `logits.shape` with `axis` removed.
+ */
+declare const categorical: OwnedFunction<(key: ArrayLike, logits: ArrayLike, args_2?: {
+  axis?: number | undefined;
+  shape?: number[] | undefined;
+  replace?: boolean | undefined;
+} | undefined) => Array>;
 /**
  * @function
  * Sample from a Cauchy distribution with location 0 and scale 1.
@@ -2645,8 +2766,31 @@ declare namespace scipy_special_d_exports {
  * The logit function, `logit(p) = log(p / (1-p))`.
  */
 declare const logit: OwnedFunction<(x: ArrayLike) => Array>;
+//#endregion
+//#region src/tracing.d.ts
+/**
+ * Start collecting kernel traces.
+ *
+ * Traces appear in developer tools under the "Performance" tab, and they are
+ * useful for measuring fine-grained kernel execution time.
+ */
+declare function startTrace(): void;
+/**
+ * Stop collecting kernel traces.
+ *
+ * Traces appear in developer tools under the "Performance" tab, and they are
+ * useful for measuring fine-grained kernel execution time.
+ */
+declare function stopTrace(): void;
+/** Check if tracing is currently enabled. */
+
 //#endregion
 //#region src/index.d.ts
+/** @namespace */
+declare const profiler: {
+  startTrace: typeof startTrace;
+  stopTrace: typeof stopTrace;
+};
 /**
  * @function
  * Compute the forward-mode Jacobian-vector product for a function.
@@ -2811,4 +2955,4 @@ declare function blockUntilReady<T extends JsTree<any>>(x: T): Promise<T>;
  */
 declare function devicePut<T extends JsTree<any>>(x: T, device?: Device): Promise<MapJsTree<T, number | boolean, Array>>;
 //#endregion
-export { Array, ClosedJaxpr, DType, type Device, Jaxpr, type JsTree, type JsTreeDef, type OwnedFunction, blockUntilReady, defaultDevice, devicePut, devices, grad, hessian, init, jacfwd, jacrev as 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 };
+export { Array, ClosedJaxpr, DType, type Device, Jaxpr, type JsTree, type JsTreeDef, type OwnedFunction, blockUntilReady, defaultDevice, devicePut, devices, grad, hessian, init, jacfwd, jacrev as jacobian, jacrev, jit, jvp, lax_d_exports as lax, linearize, makeJaxpr, nn_d_exports as nn, numpy_d_exports as numpy, profiler, random_d_exports as random, scipy_special_d_exports as scipySpecial, setDebug, tree_d_exports as tree, valueAndGrad, vjp, vmap };