@jax-js/jax 0.1.10 → 0.1.12

package/dist/index.d.cts

@@ -232,6 +232,8 @@ declare class AluExp implements FpHashable {
   static cast(dtype: DType, a: AluExp): AluExp;
   static bitcast(dtype: DType, a: AluExp): AluExp;
   static threefry2x32(k0: AluExp, k1: AluExp, c0: AluExp, c1: AluExp, mode?: "xor" | 0 | 1): AluExp;
+  static bitCombine(a: AluExp, b: AluExp, mode: "and" | "or" | "xor"): AluExp;
+  static bitShift(a: AluExp, b: AluExp, mode: "shl" | "shr"): AluExp;
   static cmplt(a: AluExp, b: AluExp): AluExp;
   static cmpne(a: AluExp, b: AluExp): AluExp;
   static where(cond: AluExp, a: AluExp, b: AluExp): AluExp;
@@ -323,6 +325,11 @@ declare enum AluOp {
   Reciprocal = "Reciprocal",
   Cast = "Cast",
   Bitcast = "Bitcast",
+  BitCombine = "BitCombine",
+  // arg = 'or' | 'and' | 'xor'
+  BitInvert = "BitInvert",
+  BitShift = "BitShift",
+  // arg = 'shl' | 'shr'
   Cmplt = "Cmplt",
   Cmpne = "Cmpne",
   Where = "Where",
@@ -546,6 +553,11 @@ declare class Executable<T = any> {
   source: Kernel | Routine, /** Extra data specific to the backend running this executable. */
   data: T);
 }
+/**
+ * If the WebGPU backend has been initialized, return the `GPUDevice` that this
+ * backend runs on. This is useful for sharing buffers.
+ */
+declare function getWebGPUDevice(): GPUDevice;
 declare namespace tree_d_exports {
   export { JsTree, JsTreeDef, MapJsTree, NodeType, dispose, flatten, leaves, map, ref, structure, unflatten };
 }
@@ -719,6 +731,8 @@ declare enum Primitive {
   // uses sign of numerator, C-style, matches JS but not Python
   Min = "min",
   Max = "max",
+  BitCombine = "bit_combine",
+  BitShift = "bit_shift",
   Neg = "neg",
   Reciprocal = "reciprocal",
   Floor = "floor",
@@ -767,6 +781,12 @@ declare enum Primitive {
   Jit = "jit",
 }
 interface PrimitiveParamsImpl extends Record<Primitive, Record<string, any>> {
+  [Primitive.BitCombine]: {
+    op: "and" | "or" | "xor";
+  };
+  [Primitive.BitShift]: {
+    op: "shl" | "shr";
+  };
   [Primitive.Cast]: {
     dtype: DType;
   };
@@ -1194,6 +1214,19 @@ declare class Array extends Tracer {
    * recommended for performance reasons, as it will block rendering.
    */
   dataSync(): DataArray;
+  /**
+   * Return this array as a WebGPU buffer (with `STORAGE | COPY_SRC`).
+   *
+   * Only available on the WebGPU backend. The array's memory is still managed
+   * by jax-js, and it will be freed when the buffer is no longer in use. You
+   * _should not_ mutate the buffer's contents.
+   *
+   * Note that the GPU buffer may be slightly larger than the array's size; it
+   * will always be aligned to 4 bytes.
+   */
+  gpuBuffer(): Promise<GPUBuffer>;
+  /** Synchronous version of `Array.gpuBuffer()`. */
+  gpuBufferSync(): GPUBuffer;
   /**
    * Convert this array into a JavaScript object.
    *
@@ -1571,7 +1604,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, cross$1 as cross, 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, vectorNorm };
 }
 /**
  * Compute the Cholesky decomposition of a (batched) positive-definite matrix.
@@ -1626,6 +1659,24 @@ declare function slogdet(a: ArrayLike): [Array, Array];
  * @returns Solution `x` of shape `(..., N)` or `(..., N, M)`.
  */
 declare function solve(a: ArrayLike, b: ArrayLike): Array;
+/**
+ * Compute the vector norm of an array.
+ *
+ * @param x - Input array.
+ * @param ord - Order of the norm (default 2). Supports `Infinity`, `-Infinity`, `0`, or any real number.
+ * @param axis - Axis/axes to reduce over (default: all axes).
+ * @param keepdims - Whether to keep reduced dimensions as size 1.
+ * @returns The norm of `x`, reduced over the given axes.
+ */
+declare function vectorNorm(x: ArrayLike, {
+  ord,
+  axis,
+  keepdims
+}?: {
+  ord?: number;
+  axis?: number | number[] | null;
+  keepdims?: boolean;
+}): Array;
 //#endregion
 //#region src/library/numpy/dtype-info.d.ts
 /** @inline */
@@ -1679,7 +1730,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, 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 };
+  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, bitwiseAnd, invert as bitwiseInvert, leftShift as bitwiseLeftShift, invert as bitwiseNot, bitwiseOr, rightShift as bitwiseRightShift, bitwiseXor, 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, invert, isfinite, isinf, isnan, isneginf, isposinf, ldexp, leftShift, 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, rightShift, 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;
@@ -1747,6 +1798,18 @@ declare function logicalOr(x: ArrayLike, y: ArrayLike): Array;
 declare function logicalXor(x: ArrayLike, y: ArrayLike): Array;
 /** Compute element-wise logical NOT. */
 declare function logicalNot(x: ArrayLike): Array;
+/** Compute element-wise bitwise AND. */
+declare function bitwiseAnd(x: ArrayLike, y: ArrayLike): Array;
+/** Compute element-wise bitwise OR. */
+declare function bitwiseOr(x: ArrayLike, y: ArrayLike): Array;
+/** Compute element-wise bitwise XOR. */
+declare function bitwiseXor(x: ArrayLike, y: ArrayLike): Array;
+/** Compute element-wise bitwise NOT (inversion). */
+declare function invert(x: ArrayLike): Array;
+/** Compute element-wise left bit shift. */
+declare function leftShift(x: ArrayLike, y: ArrayLike): Array;
+/** Compute element-wise right bit shift. */
+declare function rightShift(x: ArrayLike, y: ArrayLike): Array;
 /** @function Element-wise ternary operator, evaluates to `x` if cond else `y`. */
 declare const where: (cond: ArrayLike, x: ArrayLike, y: ArrayLike) => Array;
 /**
@@ -2659,7 +2722,7 @@ declare function dotProductAttention(query: ArrayLike, key: ArrayLike, value: Ar
   localWindowSize?: number | [number, number];
 }): Array;
 declare namespace random_d_exports {
-  export { bernoulli, bits, categorical, cauchy, exponential, gumbel, key, laplace, multivariateNormal, normal, split, uniform };
+  export { ball, bernoulli, bits, categorical, cauchy, choice, doubleSidedMaxwell, exponential, geometric, gumbel, key, laplace, logistic, lognormal, maxwell, multivariateNormal, normal, pareto, permutation, rademacher, randint, rayleigh, split, triangular, uniform, weibullMin };
 }
 /** Create a pseudo-random number generator (PRNG) key from 32-bit integer seed. */
 declare function key(seed: ArrayLike): Array;
@@ -2675,6 +2738,16 @@ declare const uniform: OwnedFunction<(key: ArrayLike, shape?: number[] | undefin
   minval?: number | undefined;
   maxval?: number | undefined;
 } | undefined) => Array>;
+/**
+ * @function
+ * Sample points uniformly from the Euclidean unit ball in `d` dimensions.
+ *
+ * Only the Euclidean `p=2` case is currently supported.
+ */
+declare const ball: OwnedFunction<(key: ArrayLike, d: number, args_2?: {
+  p?: number | undefined;
+  shape?: number[] | undefined;
+} | undefined) => Array>;
 /**
  * Sample Bernoulli random variables with given mean (0,1 categorical).
  *
@@ -2715,11 +2788,42 @@ declare const categorical: OwnedFunction<(key: ArrayLike, logits: ArrayLike, arg
  * Uses inverse transform sampling: `x = tan(π * (u - 0.5))` where u ~ Uniform(0, 1).
  */
 declare const cauchy: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined) => Array>;
+/**
+ * Sample from a population with optional replacement and optional probabilities.
+ *
+ * This implements the common JAX-compatible cases: integer populations and
+ * array populations along `axis`. Probabilities `p`, if provided, are sampled
+ * via `categorical(log(p))`.
+ */
+declare function choice(key: Array, a: number | ArrayLike, {
+  shape,
+  replace,
+  p,
+  axis
+}?: {
+  shape?: number[];
+  replace?: boolean;
+  p?: ArrayLike;
+  axis?: number;
+}): Array;
+/**
+ * @function
+ * Sample double-sided Maxwell random values with the provided location and scale.
+ */
+declare const doubleSidedMaxwell: OwnedFunction<(key: ArrayLike, loc: ArrayLike, scale: ArrayLike, shape?: number[] | undefined) => Array>;
 /**
  * @function
  * Sample exponential random values according to `p(x) = exp(-x)`.
  */
 declare const exponential: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined) => Array>;
+/**
+ * @function
+ * Sample geometric random values: the number of trials until first success.
+ */
+declare const geometric: OwnedFunction<(key: ArrayLike, p: ArrayLike, args_2?: {
+  shape?: number[] | undefined;
+  dtype?: DType | undefined;
+} | undefined) => Array>;
 /**
  * @function
  * Sample from a Gumbel distribution with location 0 and scale 1.
@@ -2735,6 +2839,23 @@ declare const gumbel: OwnedFunction<(key: ArrayLike, shape?: number[] | undefine
  * Inverting: `x = -sign(u - 0.5) * log(1 - 2 * |u - 0.5|)`.
  */
 declare const laplace: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined) => Array>;
+/**
+ * @function
+ * Sample from a logistic distribution with location 0 and scale 1.
+ *
+ * Uses inverse transform sampling: `x = log(u) - log(1-u)`.
+ */
+declare const logistic: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined) => Array>;
+/**
+ * @function
+ * Sample log-normal random values: `exp(sigma * normal(key, shape))`.
+ */
+declare const lognormal: OwnedFunction<(key: ArrayLike, sigma?: ArrayLike | undefined, shape?: number[] | undefined) => Array>;
+/**
+ * @function
+ * Sample Maxwell-distributed random values.
+ */
+declare const maxwell: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined) => Array>;
 /**
  * @function
  * Sample multivariate normal random values with given mean and covariance.
@@ -2761,6 +2882,53 @@ declare const multivariateNormal: OwnedFunction<(key: ArrayLike, mean: ArrayLike
  * bitwise identical to JAX.
  */
 declare const normal: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined) => Array>;
+/**
+ * @function
+ * Sample from a Pareto distribution with shape parameter `b` and support [1, ∞).
+ */
+declare const pareto: OwnedFunction<(key: ArrayLike, b: ArrayLike, shape?: number[] | undefined) => Array>;
+/**
+ * Return a random permutation of an integer range or of an array along `axis`.
+ */
+declare function permutation(key: Array, x: number | ArrayLike, axis?: number): Array;
+/**
+ * @function
+ * Sample Rademacher random values, uniformly from {-1, 1}.
+ */
+declare const rademacher: OwnedFunction<(key: ArrayLike, args_1?: {
+  shape?: number[] | undefined;
+  dtype?: DType | undefined;
+} | undefined) => Array>;
+/**
+ * @function
+ * Sample integer values uniformly from `[minval, maxval)`.
+ *
+ * This uses modulo reduction of uniform 32-bit random bits. For ranges that do
+ * not divide 2^32, this introduces a very small modulo bias.
+ */
+declare const randint: OwnedFunction<(key: ArrayLike, args_1: {
+  minval: number;
+  maxval: number;
+  shape?: number[] | undefined;
+  dtype?: DType | undefined;
+}) => Array>;
+/**
+ * @function
+ * Sample Rayleigh random values with the provided scale parameter.
+ */
+declare const rayleigh: OwnedFunction<(key: ArrayLike, scale?: ArrayLike | undefined, shape?: number[] | undefined) => Array>;
+/**
+ * @function
+ * Sample triangular random values on `[left, right]` with the given mode.
+ */
+declare const triangular: OwnedFunction<(key: ArrayLike, left: ArrayLike, mode: ArrayLike, right: ArrayLike, shape?: number[] | undefined) => Array>;
+/**
+ * @function
+ * Sample Weibull minimum random values.
+ *
+ * Uses `scale * exponential(key) ** (1 / concentration)`.
+ */
+declare const weibullMin: OwnedFunction<(key: ArrayLike, scale: ArrayLike, concentration: ArrayLike, shape?: number[] | undefined) => Array>;
 declare namespace scipy_special_d_exports {
   export { erf, erfc, logSoftmax, logit, logsumexp, softmax };
 }
@@ -2958,4 +3126,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, profiler, 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, getWebGPUDevice, 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 };

package/dist/index.d.ts

@@ -229,6 +229,8 @@ declare class AluExp implements FpHashable {
   static cast(dtype: DType, a: AluExp): AluExp;
   static bitcast(dtype: DType, a: AluExp): AluExp;
   static threefry2x32(k0: AluExp, k1: AluExp, c0: AluExp, c1: AluExp, mode?: "xor" | 0 | 1): AluExp;
+  static bitCombine(a: AluExp, b: AluExp, mode: "and" | "or" | "xor"): AluExp;
+  static bitShift(a: AluExp, b: AluExp, mode: "shl" | "shr"): AluExp;
   static cmplt(a: AluExp, b: AluExp): AluExp;
   static cmpne(a: AluExp, b: AluExp): AluExp;
   static where(cond: AluExp, a: AluExp, b: AluExp): AluExp;
@@ -320,6 +322,11 @@ declare enum AluOp {
   Reciprocal = "Reciprocal",
   Cast = "Cast",
   Bitcast = "Bitcast",
+  BitCombine = "BitCombine",
+  // arg = 'or' | 'and' | 'xor'
+  BitInvert = "BitInvert",
+  BitShift = "BitShift",
+  // arg = 'shl' | 'shr'
   Cmplt = "Cmplt",
   Cmpne = "Cmpne",
   Where = "Where",
@@ -543,6 +550,11 @@ declare class Executable<T = any> {
   source: Kernel | Routine, /** Extra data specific to the backend running this executable. */
   data: T);
 }
+/**
+ * If the WebGPU backend has been initialized, return the `GPUDevice` that this
+ * backend runs on. This is useful for sharing buffers.
+ */
+declare function getWebGPUDevice(): GPUDevice;
 declare namespace tree_d_exports {
   export { JsTree, JsTreeDef, MapJsTree, NodeType, dispose, flatten, leaves, map, ref, structure, unflatten };
 }
@@ -716,6 +728,8 @@ declare enum Primitive {
   // uses sign of numerator, C-style, matches JS but not Python
   Min = "min",
   Max = "max",
+  BitCombine = "bit_combine",
+  BitShift = "bit_shift",
   Neg = "neg",
   Reciprocal = "reciprocal",
   Floor = "floor",
@@ -764,6 +778,12 @@ declare enum Primitive {
   Jit = "jit",
 }
 interface PrimitiveParamsImpl extends Record<Primitive, Record<string, any>> {
+  [Primitive.BitCombine]: {
+    op: "and" | "or" | "xor";
+  };
+  [Primitive.BitShift]: {
+    op: "shl" | "shr";
+  };
   [Primitive.Cast]: {
     dtype: DType;
   };
@@ -1191,6 +1211,19 @@ declare class Array extends Tracer {
    * recommended for performance reasons, as it will block rendering.
    */
   dataSync(): DataArray;
+  /**
+   * Return this array as a WebGPU buffer (with `STORAGE | COPY_SRC`).
+   *
+   * Only available on the WebGPU backend. The array's memory is still managed
+   * by jax-js, and it will be freed when the buffer is no longer in use. You
+   * _should not_ mutate the buffer's contents.
+   *
+   * Note that the GPU buffer may be slightly larger than the array's size; it
+   * will always be aligned to 4 bytes.
+   */
+  gpuBuffer(): Promise<GPUBuffer>;
+  /** Synchronous version of `Array.gpuBuffer()`. */
+  gpuBufferSync(): GPUBuffer;
   /**
    * Convert this array into a JavaScript object.
    *
@@ -1568,7 +1601,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, cross$1 as cross, 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, vectorNorm };
 }
 /**
  * Compute the Cholesky decomposition of a (batched) positive-definite matrix.
@@ -1623,6 +1656,24 @@ declare function slogdet(a: ArrayLike): [Array, Array];
  * @returns Solution `x` of shape `(..., N)` or `(..., N, M)`.
  */
 declare function solve(a: ArrayLike, b: ArrayLike): Array;
+/**
+ * Compute the vector norm of an array.
+ *
+ * @param x - Input array.
+ * @param ord - Order of the norm (default 2). Supports `Infinity`, `-Infinity`, `0`, or any real number.
+ * @param axis - Axis/axes to reduce over (default: all axes).
+ * @param keepdims - Whether to keep reduced dimensions as size 1.
+ * @returns The norm of `x`, reduced over the given axes.
+ */
+declare function vectorNorm(x: ArrayLike, {
+  ord,
+  axis,
+  keepdims
+}?: {
+  ord?: number;
+  axis?: number | number[] | null;
+  keepdims?: boolean;
+}): Array;
 //#endregion
 //#region src/library/numpy/dtype-info.d.ts
 /** @inline */
@@ -1676,7 +1727,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, 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 };
+  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, bitwiseAnd, invert as bitwiseInvert, leftShift as bitwiseLeftShift, invert as bitwiseNot, bitwiseOr, rightShift as bitwiseRightShift, bitwiseXor, 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, invert, isfinite, isinf, isnan, isneginf, isposinf, ldexp, leftShift, 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, rightShift, 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;
@@ -1744,6 +1795,18 @@ declare function logicalOr(x: ArrayLike, y: ArrayLike): Array;
 declare function logicalXor(x: ArrayLike, y: ArrayLike): Array;
 /** Compute element-wise logical NOT. */
 declare function logicalNot(x: ArrayLike): Array;
+/** Compute element-wise bitwise AND. */
+declare function bitwiseAnd(x: ArrayLike, y: ArrayLike): Array;
+/** Compute element-wise bitwise OR. */
+declare function bitwiseOr(x: ArrayLike, y: ArrayLike): Array;
+/** Compute element-wise bitwise XOR. */
+declare function bitwiseXor(x: ArrayLike, y: ArrayLike): Array;
+/** Compute element-wise bitwise NOT (inversion). */
+declare function invert(x: ArrayLike): Array;
+/** Compute element-wise left bit shift. */
+declare function leftShift(x: ArrayLike, y: ArrayLike): Array;
+/** Compute element-wise right bit shift. */
+declare function rightShift(x: ArrayLike, y: ArrayLike): Array;
 /** @function Element-wise ternary operator, evaluates to `x` if cond else `y`. */
 declare const where: (cond: ArrayLike, x: ArrayLike, y: ArrayLike) => Array;
 /**
@@ -2656,7 +2719,7 @@ declare function dotProductAttention(query: ArrayLike, key: ArrayLike, value: Ar
   localWindowSize?: number | [number, number];
 }): Array;
 declare namespace random_d_exports {
-  export { bernoulli, bits, categorical, cauchy, exponential, gumbel, key, laplace, multivariateNormal, normal, split, uniform };
+  export { ball, bernoulli, bits, categorical, cauchy, choice, doubleSidedMaxwell, exponential, geometric, gumbel, key, laplace, logistic, lognormal, maxwell, multivariateNormal, normal, pareto, permutation, rademacher, randint, rayleigh, split, triangular, uniform, weibullMin };
 }
 /** Create a pseudo-random number generator (PRNG) key from 32-bit integer seed. */
 declare function key(seed: ArrayLike): Array;
@@ -2672,6 +2735,16 @@ declare const uniform: OwnedFunction<(key: ArrayLike, shape?: number[] | undefin
   minval?: number | undefined;
   maxval?: number | undefined;
 } | undefined) => Array>;
+/**
+ * @function
+ * Sample points uniformly from the Euclidean unit ball in `d` dimensions.
+ *
+ * Only the Euclidean `p=2` case is currently supported.
+ */
+declare const ball: OwnedFunction<(key: ArrayLike, d: number, args_2?: {
+  p?: number | undefined;
+  shape?: number[] | undefined;
+} | undefined) => Array>;
 /**
  * Sample Bernoulli random variables with given mean (0,1 categorical).
  *
@@ -2712,11 +2785,42 @@ declare const categorical: OwnedFunction<(key: ArrayLike, logits: ArrayLike, arg
  * Uses inverse transform sampling: `x = tan(π * (u - 0.5))` where u ~ Uniform(0, 1).
  */
 declare const cauchy: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined) => Array>;
+/**
+ * Sample from a population with optional replacement and optional probabilities.
+ *
+ * This implements the common JAX-compatible cases: integer populations and
+ * array populations along `axis`. Probabilities `p`, if provided, are sampled
+ * via `categorical(log(p))`.
+ */
+declare function choice(key: Array, a: number | ArrayLike, {
+  shape,
+  replace,
+  p,
+  axis
+}?: {
+  shape?: number[];
+  replace?: boolean;
+  p?: ArrayLike;
+  axis?: number;
+}): Array;
+/**
+ * @function
+ * Sample double-sided Maxwell random values with the provided location and scale.
+ */
+declare const doubleSidedMaxwell: OwnedFunction<(key: ArrayLike, loc: ArrayLike, scale: ArrayLike, shape?: number[] | undefined) => Array>;
 /**
  * @function
  * Sample exponential random values according to `p(x) = exp(-x)`.
  */
 declare const exponential: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined) => Array>;
+/**
+ * @function
+ * Sample geometric random values: the number of trials until first success.
+ */
+declare const geometric: OwnedFunction<(key: ArrayLike, p: ArrayLike, args_2?: {
+  shape?: number[] | undefined;
+  dtype?: DType | undefined;
+} | undefined) => Array>;
 /**
  * @function
  * Sample from a Gumbel distribution with location 0 and scale 1.
@@ -2732,6 +2836,23 @@ declare const gumbel: OwnedFunction<(key: ArrayLike, shape?: number[] | undefine
  * Inverting: `x = -sign(u - 0.5) * log(1 - 2 * |u - 0.5|)`.
  */
 declare const laplace: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined) => Array>;
+/**
+ * @function
+ * Sample from a logistic distribution with location 0 and scale 1.
+ *
+ * Uses inverse transform sampling: `x = log(u) - log(1-u)`.
+ */
+declare const logistic: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined) => Array>;
+/**
+ * @function
+ * Sample log-normal random values: `exp(sigma * normal(key, shape))`.
+ */
+declare const lognormal: OwnedFunction<(key: ArrayLike, sigma?: ArrayLike | undefined, shape?: number[] | undefined) => Array>;
+/**
+ * @function
+ * Sample Maxwell-distributed random values.
+ */
+declare const maxwell: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined) => Array>;
 /**
  * @function
  * Sample multivariate normal random values with given mean and covariance.
@@ -2758,6 +2879,53 @@ declare const multivariateNormal: OwnedFunction<(key: ArrayLike, mean: ArrayLike
  * bitwise identical to JAX.
  */
 declare const normal: OwnedFunction<(key: ArrayLike, shape?: number[] | undefined) => Array>;
+/**
+ * @function
+ * Sample from a Pareto distribution with shape parameter `b` and support [1, ∞).
+ */
+declare const pareto: OwnedFunction<(key: ArrayLike, b: ArrayLike, shape?: number[] | undefined) => Array>;
+/**
+ * Return a random permutation of an integer range or of an array along `axis`.
+ */
+declare function permutation(key: Array, x: number | ArrayLike, axis?: number): Array;
+/**
+ * @function
+ * Sample Rademacher random values, uniformly from {-1, 1}.
+ */
+declare const rademacher: OwnedFunction<(key: ArrayLike, args_1?: {
+  shape?: number[] | undefined;
+  dtype?: DType | undefined;
+} | undefined) => Array>;
+/**
+ * @function
+ * Sample integer values uniformly from `[minval, maxval)`.
+ *
+ * This uses modulo reduction of uniform 32-bit random bits. For ranges that do
+ * not divide 2^32, this introduces a very small modulo bias.
+ */
+declare const randint: OwnedFunction<(key: ArrayLike, args_1: {
+  minval: number;
+  maxval: number;
+  shape?: number[] | undefined;
+  dtype?: DType | undefined;
+}) => Array>;
+/**
+ * @function
+ * Sample Rayleigh random values with the provided scale parameter.
+ */
+declare const rayleigh: OwnedFunction<(key: ArrayLike, scale?: ArrayLike | undefined, shape?: number[] | undefined) => Array>;
+/**
+ * @function
+ * Sample triangular random values on `[left, right]` with the given mode.
+ */
+declare const triangular: OwnedFunction<(key: ArrayLike, left: ArrayLike, mode: ArrayLike, right: ArrayLike, shape?: number[] | undefined) => Array>;
+/**
+ * @function
+ * Sample Weibull minimum random values.
+ *
+ * Uses `scale * exponential(key) ** (1 / concentration)`.
+ */
+declare const weibullMin: OwnedFunction<(key: ArrayLike, scale: ArrayLike, concentration: ArrayLike, shape?: number[] | undefined) => Array>;
 declare namespace scipy_special_d_exports {
   export { erf, erfc, logSoftmax, logit, logsumexp, softmax };
 }
@@ -2955,4 +3123,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, profiler, 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, getWebGPUDevice, 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 };