@jax-js/jax 0.1.1 → 0.1.3

package/dist/index.d.ts

@@ -224,6 +224,8 @@ declare class AluExp implements FpHashable {
   static erf(a: AluExp): AluExp;
   static erfc(a: AluExp): AluExp;
   static sqrt(a: AluExp): AluExp;
+  static floor(a: AluExp): AluExp;
+  static ceil(a: AluExp): AluExp;
   static reciprocal(a: AluExp): AluExp;
   static cast(dtype: DType, a: AluExp): AluExp;
   static bitcast(dtype: DType, a: AluExp): AluExp;
@@ -249,7 +251,7 @@ declare class AluExp implements FpHashable {
   /** Substitute variables in this AluExp to values. */
   substitute(variables: Record<string, AluExp>): AluExp;
   /** Reindex gid values in this expression as needed. */
-  reindexGids(gidMap: Map<number, number>): AluExp;
+  reindexGids(newGids: number[]): AluExp;
   get min(): number;
   get max(): number;
   /** Largest known integer that divides self. */
@@ -314,6 +316,8 @@ declare enum AluOp {
   Erf = "Erf",
   Erfc = "Erfc",
   Sqrt = "Sqrt",
+  Floor = "Floor",
+  Ceil = "Ceil",
   Reciprocal = "Reciprocal",
   Cast = "Cast",
   Bitcast = "Bitcast",
@@ -456,679 +460,85 @@ declare class Executable<T = any> {
   constructor(kernel: Kernel, /** Extra data specific to the backend running this kernel. */
   data: T);
 }
-declare namespace tree_d_exports {
-  export { JsTree, JsTreeDef, MapJsTree, NodeType, dispose, flatten, leaves, map, ref, structure, unflatten };
-}
-declare enum NodeType {
-  Array = "Array",
-  Object = "Object",
-  Leaf = "Leaf",
-}
-/** Analog to the JAX "pytree" object, but for JavaScript. */
-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 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. */
-declare class JsTreeDef {
-  readonly nodeType: NodeType;
-  readonly nodeMetadata: any;
-  readonly childTreedefs: JsTreeDef[];
-  static leaf: 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. */
-  equals(other: JsTreeDef): boolean;
-}
-/** Flatten a structured object, returning the tree definition. */
-declare function flatten<T>(tree: JsTree<T>): [T[], JsTreeDef];
-/** Get the leaves of a tree. */
-declare function leaves<T>(tree: JsTree<T>): T[];
-/** Get the treedef for a tree. */
-declare function structure<T>(tree: JsTree<T>): JsTreeDef;
-/** Reconstruct a structured object from the flattened representation. */
-declare function unflatten<T>(treedef: JsTreeDef, leaves: Iterable<T>): JsTree<T>;
-/** Maps a multi-input function over pytree args to produce a new pytree. */
-declare function map<T, U, Tree extends JsTree<T>>(fn: (...args: T[]) => U, tree: Tree, ...rest: Tree[]): MapJsTree<Tree, T, U>;
-/** Take a reference of every array in a tree. */
-declare function ref<Tree extends JsTree<Array>>(tree: Tree): Tree;
-/** Dispose every array in a tree. */
-declare function dispose<Tree extends JsTree<Array>>(tree: Tree | null | undefined): void;
-//#endregion
-//#region src/frontend/convolution.d.ts
-/** Definition of a general dilated convolution. Should be valid on creation. */
-interface ConvParams {
-  strides: number[];
-  padding: [number, number][];
-  lhsDilation: number[];
-  rhsDilation: number[];
+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, ceil, clip, columnStack, concatenate, cos, cosh, cumsum, cumulativeSum, deg2rad, degrees, diag, diagonal, divide, dot$1 as dot, dstack, e, einsum, equal, eulerGamma, exp, exp2, expm1, eye, flip, fliplr, flipud, float16, float32, float64, floor, fmod, frexp, full, fullLike, greater, greaterEqual, hamming, hann, heaviside, hstack, hypot, identity$1 as identity, inf, inner, int32, isfinite, isinf, isnan, isneginf, isposinf, ldexp, 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, positive, pow, power, prod, promoteTypes, ptp, rad2deg, radians, ravel, reciprocal, remainder, repeat, reshape, shape$1 as shape, sign, sin, sinh, size, sqrt, square, squeeze, stack, std, subtract, sum, tan, tanh, tensordot, tile, trace, 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...` */
+declare const eulerGamma = 0.5772156649015329;
+/** Positive infinity. */
+declare const inf: number;
+/** Floating-point representation of NaN. */
+declare const nan: number;
+/** This is Pi, `π = 3.14159265358979...` */
+declare const pi: number;
+/** @function Element-wise addition, with broadcasting. */
+declare const add: (x: ArrayLike, y: ArrayLike) => Array;
+/** @function Element-wise multiplication, with broadcasting. */
+declare const multiply: (x: ArrayLike, y: ArrayLike) => Array;
+/** @function Numerical negative of every element of an array. */
+declare const negative: (x: ArrayLike) => Array;
+/** @function Calculate element-wise reciprocal of the input. This is `1/x`. */
+declare const reciprocal: (x: ArrayLike) => Array;
+/** @function Round input down to the nearest integer. */
+declare const floor: (x: ArrayLike) => Array;
+/** @function Round input up to the nearest integer. */
+declare const ceil: (x: ArrayLike) => Array;
+/** @function Element-wise sine function (takes radians). */
+declare const sin: (x: ArrayLike) => Array;
+/** @function Element-wise cosine function (takes radians). */
+declare const cos: (x: ArrayLike) => Array;
+/** @function Element-wise inverse sine function (inverse of sin). */
+declare const asin: (x: ArrayLike) => Array;
+/** @function Element-wise inverse tangent function (inverse of tan). */
+declare const atan: (x: ArrayLike) => Array;
+/** @function Calculate the exponential of all elements in the input array. */
+declare const exp: (x: ArrayLike) => Array;
+/** @function Calculate the natural logarithm of all elements in the input array. */
+declare const log: (x: ArrayLike) => Array;
+/** @function Calculate the square root of all elements in the input array. */
+declare const sqrt: (x: ArrayLike) => Array;
+/** @function Return element-wise minimum of the input arrays. */
+declare const minimum: (x: ArrayLike, y: ArrayLike) => Array;
+/** @function Return element-wise maximum of the input arrays. */
+declare const maximum: (x: ArrayLike, y: ArrayLike) => Array;
+/** @function Compare two arrays element-wise. */
+declare const greater: (x: ArrayLike, y: ArrayLike) => Array;
+/** @function Compare two arrays element-wise. */
+declare const less: (x: ArrayLike, y: ArrayLike) => Array;
+/** @function Compare two arrays element-wise. */
+declare const equal: (x: ArrayLike, y: ArrayLike) => Array;
+/** @function Compare two arrays element-wise. */
+declare const notEqual: (x: ArrayLike, y: ArrayLike) => Array;
+/** @function Compare two arrays element-wise. */
+declare const greaterEqual: (x: ArrayLike, y: ArrayLike) => Array;
+/** @function Compare two arrays element-wise. */
+declare const lessEqual: (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;
 /**
- * Check that the shapes and parameters passed to convolution are valid.
- *
- * If the check succeeds, returns the output shape.
+ * @function
+ * Permute the dimensions of an array. Defaults to reversing the axis order.
  */
-
-//#endregion
-//#region src/frontend/core.d.ts
+declare const transpose: (x: ArrayLike, perm?: number[]) => Array;
 /**
- * Frontend primitive operations, which are lowered into Kernel objects before
- * being dispatched to the backend.
- *
- * Any operation between arrays can be described in these parts. This is also
- * the set of primitives that can occur in Jaxpr programs, and the level at
- * which transformations like vmap, grad, and jvp occur. They are loosely based
- * on [XLA](https://openxla.org/xla/operation_semantics).
+ * @function
+ * Give a new shape to an array without changing its data.
  *
- * All n-ary operations support broadcasting, with NumPy semantics.
+ * One shape dimension can be -1. In this case, the value is inferred from the
+ * length of the array and remaining dimensions.
  */
-declare enum Primitive {
-  Add = "add",
-  Mul = "mul",
-  Idiv = "idiv",
-  Neg = "neg",
-  Reciprocal = "reciprocal",
-  StopGradient = "stop_gradient",
-  Cast = "cast",
-  Bitcast = "bitcast",
-  RandomBits = "random_bits",
-  Sin = "sin",
-  Cos = "cos",
-  Asin = "asin",
-  Atan = "atan",
-  Exp = "exp",
-  Log = "log",
-  Erf = "erf",
-  Erfc = "erfc",
-  Sqrt = "sqrt",
-  Min = "min",
-  Max = "max",
-  Reduce = "reduce",
-  Dot = "dot",
-  // sum(x*y, axis=-1)
-  Conv = "conv",
-  // see lax.conv_general_dilated
-  Pool = "pool",
-  PoolTranspose = "pool_transpose",
-  Compare = "compare",
-  Where = "where",
-  Transpose = "transpose",
-  Broadcast = "broadcast",
-  Reshape = "reshape",
-  Flip = "flip",
-  Shrink = "shrink",
-  Pad = "pad",
-  Gather = "gather",
-  JitCall = "jit_call",
-}
-interface PrimitiveParamsImpl extends Record<Primitive, Record<string, any>> {
-  [Primitive.Cast]: {
-    dtype: DType;
-  };
-  [Primitive.Bitcast]: {
-    dtype: DType;
-  };
-  [Primitive.Reduce]: {
-    op: AluOp;
-    axis: number[];
-  };
-  [Primitive.Conv]: ConvParams;
-  [Primitive.Pool]: {
-    window: number[];
-    strides: number[];
-  };
-  [Primitive.PoolTranspose]: {
-    inShape: number[];
-    window: number[];
-    strides: number[];
-  };
-  [Primitive.Compare]: {
-    op: CompareOp;
-  };
-  [Primitive.Transpose]: {
-    perm: number[];
-  };
-  [Primitive.Broadcast]: {
-    shape: number[];
-    axis: number[];
-  };
-  [Primitive.RandomBits]: {
-    shape: number[];
-    mode: "xor" | 0 | 1;
-  };
-  [Primitive.Reshape]: {
-    shape: number[];
-  };
-  [Primitive.Flip]: {
-    axis: number[];
-  };
-  [Primitive.Shrink]: {
-    slice: Pair[];
-  };
-  [Primitive.Pad]: {
-    width: Pair[];
-  };
-  [Primitive.Gather]: {
-    axis: number[];
-    outDim: number;
-  };
-  [Primitive.JitCall]: {
-    name: string;
-    jaxpr: Jaxpr;
-    numConsts: number;
-  };
-}
-/** Type of parameters taken by each primitive. */
-type PrimitiveParams<T extends Primitive> = T extends keyof PrimitiveParamsImpl ? PrimitiveParamsImpl[T] : Record<string, never>;
-declare enum CompareOp {
-  Less = "less",
-  Equal = "equal",
-  NotEqual = "not_equal",
-  LessEqual = "less_equal",
-}
-/** @inline */
-type Axis = number | number[] | null;
-/** @inline */
-type ReduceOpts = {
-  keepdims?: boolean;
-};
-type MainTrace = {
-  level: number;
-  traceType: new (main: MainTrace) => Trace;
-  globalData: any | null;
-};
+declare const reshape: (x: ArrayLike, shape: number[]) => Array;
 /**
- * Push an interpreter onto the trace stack. Use this like:
- * `using main = newMain(...);`
- */
-
-type TracerValue = Tracer | number | boolean;
-declare abstract class Trace {
-  readonly main: MainTrace;
-  constructor(main: MainTrace);
-  abstract pure(val: TracerValue): Tracer;
-  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;
-  constructor(trace: Trace);
-  abstract get aval(): AbstractValue;
-  abstract toString(): string;
-  /**
-   * Access an array by reference, incrementing the reference count.
-   *
-   * jax-js handles freeing arrays by using "move" semantics, like in Rust/C++.
-   * Whenever you pass an array into a function, that function should consume
-   * the array, and it will no longer be usable. For example, if you had:
-   *
-   * ```
-   * const x = np.array([1, 2, 3]);
-   * const y = np.add(x, x);
-   * ```
-   *
-   * The second line does not work because the first parameter consumes `x`, and
-   * then the second parameter will already have been freed / disposed.
-   *
-   * To fix this, you can write:
-   *
-   * ```
-   * const y = np.add(x.ref, x);
-   * ```
-   *
-   * Under the hood, every access to `.ref` increments the internal reference
-   * count of the array. The reference count starts at 1. When it hits 0, the
-   * memory behind the array is freed.
-   */
-  abstract get ref(): this;
-  /**
-   * Manually decrement the reference count of the array.
-   *
-   * Arrays are created with reference count 1. Whenever it is used as argument
-   * to a function or other operation, it is disposed (i.e., reference count
-   * decreases by 1) automatically. Whenever a `.ref` is created, the reference
-   * count increases.
-   *
-   * You generally don't need to call this function directly since arrays are
-   * automatically disposed after being passed into an operation. One common
-   * exception is when writing a function and ignoring one of its arguments. In
-   * that case, by convention you should dispose of that argument manually.
-   *
-   * ```
-   * function myCustomOperation(a: np.Array, b: np.Array) {
-   *   b.dispose(); // Needed to satisfy "move" rules.
-   *   return a.add(1);
-   * }
-   * ```
-   */
-  abstract dispose(): void;
-  /** The shape of the array. */
-  get shape(): number[];
-  /** The total number of elements in the array. */
-  get size(): number;
-  /** 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 */
-  fullLower(): Tracer;
-  neg(): this;
-  add(other: this | TracerValue): this;
-  mul(other: this | TracerValue): this;
-  greater(other: this | TracerValue): this;
-  less(other: this | TracerValue): this;
-  equal(other: this | TracerValue): this;
-  notEqual(other: this | TracerValue): this;
-  greaterEqual(other: this | TracerValue): this;
-  lessEqual(other: this | TracerValue): this;
-  /** Sum of the elements of the array over a given axis, or axes. */
-  sum(axis?: Axis, opts?: ReduceOpts): this;
-  /** Product of the array elements over a given axis. */
-  prod(axis?: Axis, opts?: ReduceOpts): this;
-  /** Compute the average of the array elements along the specified axis. */
-  mean(axis?: Axis, opts?: ReduceOpts): this;
-  /** Permute the dimensions of an array. Defaults to reversing the axis order. */
-  transpose(perm?: number[]): this;
-  /**
-   * Give a new shape to an array without changing its data.
-   *
-   * One shape dimension can be -1. In this case, the value is inferred from the
-   * length of the array and remaining dimensions.
-   */
-  reshape(shape: number | number[]): this;
-  /** Copy the array and cast to a specified dtype. */
-  astype(dtype: DType): this;
-  /** Subtract an array from this one. */
-  sub(other: this | TracerValue): this;
-  /** Divide an array by this one. */
-  div(other: this | TracerValue): this;
-  /** Return specified diagonals. See `numpy.diagonal` for full docs. */
-  diagonal(offset?: number, axis1?: number, axis2?: number): this;
-  /** Flatten the array without changing its data. */
-  flatten(): this;
-  /** Flatten the array without changing its data. */
-  ravel(): this;
-  /**
-   * Iterate over the first dimension of this array, returning slices.
-   *
-   * This can be used to destructure arrays. For example:
-   *
-   * ```js
-   * let x = np.array([[1, 2], [3, 4]]);
-   * let [a, b] = x;
-   * console.log(a.js()); // [1, 2]
-   * console.log(b.js()); // [3, 4]
-   * ```
-   */
-  [Symbol.iterator](): IterableIterator<this>;
-  /**
-   * Slice an array along one or more axes.
-   *
-   * This is the equivalent of slicing in Python, e.g. `x[1:3, 2, :, None]`. To
-   * mimic this in JavaScript, we would write:
-   *
-   * ```js
-   * x.slice([1, 3], 2, [], null);
-   * ```
-   *
-   * The `slice` method accepts a variable number of arguments, each of which
-   * can be a number, an empty array, a single-element array, a two-element
-   * array, or `null`. The arguments are interpreted as follows:
-   *
-   * - A number `n` means to access the `n`-th element along that axis, removing
-   *   that axis from the resulting shape.
-   * - An empty array `[]` means to keep that axis as-is, like `:` in Python.
-   * - A single-element array `[i]` means to start slicing from index `i`
-   *   (inclusive) to the end of the axis, like `x[i:]`.
-   * - A two-element array `[i, j]` means to slice from index `i` (inclusive)
-   *   to index `j` (exclusive), like `x[i:j]`.
-   * - `null` means to add a new axis at that position, like `np.newaxis`.
-   *
-   * Like in Python, negative indices are supported, which count from the end of
-   * the axis. For example, `-1` means the last element.
-   *
-   * Strided slices are not yet implemented, so you cannot write `x[::2]` or
-   * similar.
-   *
-   * Advanced indexing by integer arrays is also supported. This translates to
-   * the "gather" primitive, and it allows you to access specific elements of
-   * the array by integer indices stored in another array.
-   */
-  slice(...index: (number | [] | [number] | Pair | null | Tracer)[]): this;
-}
-declare class ShapedArray implements AbstractValue {
-  readonly shape: number[];
-  readonly dtype: DType;
-  readonly weakType: boolean;
-  constructor(shape: number[], dtype: DType, weakType: boolean);
-  static fromAval(aval: AbstractValue): ShapedArray;
-  get ndim(): number;
-  toString(): string;
-  equals(other: ShapedArray): boolean;
-}
-//#endregion
-//#region src/frontend/array.d.ts
-type ArrayLike = Array | number | boolean;
-/** Version of pureArray with fudged types. */
-
-/**
- * An executable operation that will be dispatched to the backend.
- *
- * This holds a reference to all input buffers used in the operation. After the
- * operation is dispatched, the references should be released.
- */
-declare class PendingExecute {
-  #private;
-  readonly backend: Backend;
-  readonly kernel: Kernel;
-  readonly inputs: Slot[];
-  readonly outputs: Slot[];
-  prepared: Executable | null;
-  submitted: boolean;
-  constructor(backend: Backend, kernel: Kernel, inputs: Slot[], outputs: Slot[]);
-  updateRc(delta: number): void;
-  prepare(): Promise<void>;
-  prepareSync(): void;
-  submit(): void;
-}
-/** @inline */
-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 `jax.Array` from JAX, or
- * `torch.Tensor`.
- *
- * Not to be confused with the JavaScript "Array" constructor. Avoid importing
- * this into your code's namespace if you're already using the JavaScript
- * "Array" type by name.
- */
-declare class Array extends Tracer {
-  #private;
-  id: number;
-  /**
-   * @ignore
-   * Constructs an array from source, shape and backend. Note that if the source
-   * is a backend `Slot`, this constructor _takes ownership_ of the slot. It
-   * will be freed when the array is disposed.
-   */
-  constructor(args: ArrayConstructorArgs);
-  /** @ignore */
-  get aval(): ShapedArray;
-  /** Return a simple string representation of the array's dimensions. */
-  toString(): string;
-  get device(): Device;
-  get ref(): this;
-  dispose(): void;
-  /**
-   * Convert this array into a primitive value.
-   *
-   * This only works for scalars (0-dimensional arrays). It lets you get values
-   * "out" of the JAX system. For instance, if `x = np.array(5)`, then you can
-   * evaluate `x + 1` and `x ** 2` to get `6` and `25`, respectively.
-   *
-   * This method is also called for `==` equality.
-   */
-  [Symbol.toPrimitive](): any;
-  /** Realize the array and return it as data. */
-  data(): Promise<DataArray>;
-  /**
-   * Wait for this array to finish evaluation.
-   *
-   * Operations and data loading in jax-js are lazy, so this function ensures
-   * that pending operations are dispatched and fully executed before it
-   * returns.
-   *
-   * If you are mapping from `data()` or `dataSync()`, it will also trigger
-   * dispatch of operations as well.
-   *
-   * **Note:** `jax.blockUntilReady()` is a higher-level API, it calls this
-   * asynchronously for multiple arrays.
-   */
-  blockUntilReady(): Promise<Array>;
-  /**
-   * Realize the array and return it as data. This is a sync variant and not
-   * recommended for performance reasons, as it will block rendering.
-   */
-  dataSync(): DataArray;
-  /**
-   * Convert this array into a JavaScript object.
-   *
-   * This is a blocking operation that will compile all of the shaders and wait
-   * for execution to complete, synchronously. No other JavaScript code on the
-   * site will be run during shader execution.
-   *
-   * To avoid blocking, prefer `jsAsync()` when possible.
-   */
-  js(): any;
-  /** Convert this array into a JavaScript object, asynchronously. */
-  jsAsync(): Promise<any>;
-  /**
-   * Copy an element of an array to a numeric scalar and return it.
-   *
-   * Throws an error if the array does not have a single element. The array must
-   * either be rank-0, or all dimensions of the shape are 1.
-   */
-  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 | DataArray | RecursiveArray<number> | RecursiveArray<boolean>, {
-  shape,
-  dtype,
-  device
-}?: {
-  shape?: number[];
-} & DTypeAndDevice): Array;
-/** If x is a value, lift it into an array, otherwise leave it be. */
-
-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;
-/** Return a new array of given shape and type, filled with ones. */
-declare function ones(shape: number[], {
-  dtype,
-  device
-}?: 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,
-  device
-}?: DTypeAndDevice): Array;
-/**
- * Create an identity matrix.
- *
- * If numCols is not provided, it defaults to numRows, i.e., a square identity
- * matrix with ones on the diagonal.
- */
-declare function eye(numRows: number, numCols?: number, {
-  dtype,
-  device
-}?: DTypeAndDevice): Array;
-/** Return the identity matrix, with ones on the main diagonal. */
-declare function identity$1(n: number, {
-  dtype,
-  device
-}?: DTypeAndDevice): Array;
-/**
- * Return evenly spaced values within a given interval.
- *
- * This can be called with a varying number of arguments, just like the range()
- * builtin function in Python.
- *
- * - `arange(stop)` is equivalent to `arange(0, stop, 1)`.
- * - `arange(start, stop)` is equivalent to `arange(start, stop, 1)`.
- * - `arange(start, stop, step)` creates an array starting at `start`, ending
- *   before `stop`, with a step size of `step`.
- *
- * Defaults to an integer data type. This can produce unintended results when
- * using a non-integer step, so prefer linspace() in those cases.
- */
-declare function arange(start: number, stop?: number, step?: number, {
-  dtype,
-  device
-}?: DTypeAndDevice): Array;
-/**
- * Return evenly spaced numbers over a specified interval.
- *
- * Returns _num_ evenly spaced samples, calculated over the interval
- * [`start`, `stop`]. The endpoint `stop` is included in the result by default,
- * but this is controlled by the `endpoint` parameter.
- *
- * The default data type is Float32. Use arange() for integer steps.
- */
-declare function linspace(start: number, stop: number, num?: number, endpoint?: boolean, {
-  dtype,
-  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, 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...` */
-declare const eulerGamma = 0.5772156649015329;
-/** Positive infinity. */
-declare const inf: number;
-/** Floating-point representation of NaN. */
-declare const nan: number;
-/** This is Pi, `π = 3.14159265358979...` */
-declare const pi: number;
-/** @function Element-wise addition, with broadcasting. */
-declare const add: (x: ArrayLike, y: ArrayLike) => Array;
-/** @function Element-wise multiplication, with broadcasting. */
-declare const multiply: (x: ArrayLike, y: ArrayLike) => Array;
-/** @function Numerical negative of every element of an array. */
-declare const negative: (x: ArrayLike) => Array;
-/** @function Calculate element-wise reciprocal of the input. This is `1/x`. */
-declare const reciprocal: (x: ArrayLike) => Array;
-/** @function Element-wise sine function (takes radians). */
-declare const sin: (x: ArrayLike) => Array;
-/** @function Element-wise cosine function (takes radians). */
-declare const cos: (x: ArrayLike) => Array;
-/** @function Element-wise inverse sine function (inverse of sin). */
-declare const asin: (x: ArrayLike) => Array;
-/** @function Element-wise inverse tangent function (inverse of tan). */
-declare const atan: (x: ArrayLike) => Array;
-/** @function Calculate the exponential of all elements in the input array. */
-declare const exp: (x: ArrayLike) => Array;
-/** @function Calculate the natural logarithm of all elements in the input array. */
-declare const log: (x: ArrayLike) => Array;
-/** @function Calculate the square root of all elements in the input array. */
-declare const sqrt: (x: ArrayLike) => Array;
-/** @function Return element-wise minimum of the input arrays. */
-declare const minimum: (x: ArrayLike, y: ArrayLike) => Array;
-/** @function Return element-wise maximum of the input arrays. */
-declare const maximum: (x: ArrayLike, y: ArrayLike) => Array;
-/** @function Compare two arrays element-wise. */
-declare const greater: (x: ArrayLike, y: ArrayLike) => Array;
-/** @function Compare two arrays element-wise. */
-declare const less: (x: ArrayLike, y: ArrayLike) => Array;
-/** @function Compare two arrays element-wise. */
-declare const equal: (x: ArrayLike, y: ArrayLike) => Array;
-/** @function Compare two arrays element-wise. */
-declare const notEqual: (x: ArrayLike, y: ArrayLike) => Array;
-/** @function Compare two arrays element-wise. */
-declare const greaterEqual: (x: ArrayLike, y: ArrayLike) => Array;
-/** @function Compare two arrays element-wise. */
-declare const lessEqual: (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;
-/**
- * @function
- * Permute the dimensions of an array. Defaults to reversing the axis order.
- */
-declare const transpose: (x: ArrayLike, perm?: number[]) => Array;
-/**
- * @function
- * Give a new shape to an array without changing its data.
- *
- * One shape dimension can be -1. In this case, the value is inferred from the
- * length of the array and remaining dimensions.
- */
-declare const reshape: (x: ArrayLike, shape: number[]) => Array;
-/**
- * @function
- * Move axes of an array to new positions. Other axes retain original order.
+ * @function
+ * Move axes of an array to new positions. Other axes retain original order.
  */
 declare const moveaxis: (x: ArrayLike, src: number, dst: number) => Array;
 /**
@@ -1177,6 +587,8 @@ declare function prod(a: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
 declare function min(a: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
 /** Return the maximum of array elements along a given axis. */
 declare function max(a: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
+/** Return the peak-to-peak range along a given axis (`max - min`). */
+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;
 /**
@@ -1193,6 +605,15 @@ declare function argmin(a: ArrayLike, axis?: number, opts?: ReduceOpts): Array;
  * specified axis.
  */
 declare function argmax(a: ArrayLike, axis?: number, opts?: ReduceOpts): Array;
+/**
+ * Cumulative sum of elements along an axis.
+ *
+ * Currently this function is `O(n^2)`, we'll improve this later on with a
+ * two-phase parallel reduction algorithm.
+ */
+declare function cumsum(a: ArrayLike, axis?: number): Array;
+/** @function Alternative name for `jax.numpy.cumsum()`. */
+declare const cumulativeSum: typeof cumsum;
 /** Reverse the elements in an array along the given axes. */
 declare function flip(x: ArrayLike, axis?: Axis): Array;
 /**
@@ -1242,6 +663,8 @@ declare function fliplr(x: ArrayLike): Array;
 declare const permuteDims: (x: ArrayLike, perm?: number[]) => Array;
 /** Return a 1-D flattened array containing the elements of the input. */
 declare function ravel(a: ArrayLike): Array;
+/** Remove one or more length-1 axes from an array. */
+declare function squeeze(a: ArrayLike, axis?: Axis): Array;
 /**
  * Repeat each element of an array after themselves.
  *
@@ -1250,381 +673,1078 @@ declare function ravel(a: ArrayLike): Array;
  */
 declare function repeat(a: ArrayLike, repeats: number, axis?: number): Array;
 /**
- * Construct an array by repeating A the number of times given by reps.
- *
- * If `A` is an array of shape `(d1, d2, ..., dn)` and `reps` is a sequence of
- * integers, the resulting array will have a shape of `(reps[0] * d1,
- * reps[1] * d2, ..., reps[n] * dn)`, with `A` tiled along each dimension.
+ * Construct an array by repeating A the number of times given by reps.
+ *
+ * If `A` is an array of shape `(d1, d2, ..., dn)` and `reps` is a sequence of
+ * integers, the resulting array will have a shape of `(reps[0] * d1,
+ * reps[1] * d2, ..., reps[n] * dn)`, with `A` tiled along each dimension.
+ */
+declare function tile(a: ArrayLike, reps: number | number[]): Array;
+/**
+ * Broadcast an array to a shape, with NumPy-style broadcasing rules.
+ *
+ * In other words, this lets you append axes to the left, and/or expand
+ * dimensions where the shape is 1.
+ */
+declare function broadcastTo(a: ArrayLike, shape: number[]): Array;
+/** Broadcast input shapes to a common output shape. */
+declare function broadcastShapes(...shapes: number[][]): number[];
+/** Broadcast arrays to a common shape. */
+declare function broadcastArrays(...arrays: ArrayLike[]): Array[];
+/**
+ * Return specified diagonals.
+ *
+ * If a is 2D, return the diagonal of the array with the given offset. If a is
+ * 3D or higher, compute diagonals along the two given axes (default: 0, 1).
+ *
+ * This returns a view over the existing array. The shape of the resulting array
+ * is determined by removing the two axes along which the diagonal is taken,
+ * then appending a new axis to the right with holding the diagonals.
+ */
+declare function diagonal(a: ArrayLike, offset?: number, axis1?: number, axis2?: number): Array;
+/**
+ * Extract a diagonal or construct a diagonal array.
+ *
+ * If v is a 2D array, return the k-th diagonal of v (as a view). If v is a 1D
+ * array, return a 2D array with v on the k-th diagonal.
+ */
+declare function diag(v: ArrayLike, k?: number): Array;
+/** Calculate the sum of the diagonal of an array along the given axes. */
+declare function trace(a: ArrayLike, offset?: number, axis1?: number, axis2?: number): Array;
+/** Return if two arrays are element-wise equal within a tolerance. */
+declare function allclose(actual: Parameters<typeof array>[0], expected: Parameters<typeof array>[0], options?: {
+  rtol?: number;
+  atol?: number;
+}): boolean;
+/** Matrix product of two arrays. */
+declare function matmul(x: ArrayLike, y: ArrayLike): Array;
+/** Dot product of two arrays. */
+declare function dot$1(x: ArrayLike, y: ArrayLike): Array;
+/**
+ * Compute the tensor dot product of two N-dimensional arrays.
+ *
+ * The behavior is determined by `axes`. If an integer `k`, sum over the last
+ * `k` axes of x and the first `k` axes of y. If a tuple, then the first array
+ * corresponds to the axes of x and the second to the axes of y.
+ */
+declare function tensordot(x: ArrayLike, y: ArrayLike, axes?: number | [number[], number[]]): Array;
+/**
+ * Einstein summation with string subscripts.
+ *
+ * @example
+ * ```ts
+ * import { numpy as np } from "@jax-js/jax";
+ *
+ * const a = np.ones([2, 3]);
+ * const b = np.ones([3]);
+ * np.einsum("ij,j", a, b); // Shape [2]
+ * ```
+ */
+declare function einsum(subscripts: string, ...operands: ArrayLike[]): Array;
+/**
+ * Einstein summation alternating between arrays and numeric indices.
+ *
+ * @example
+ * ```ts
+ * import { numpy as np } from "@jax-js/jax";
+ *
+ * const a = np.ones([2, 3]);
+ * const b = np.ones([3]);
+ * np.einsum(a, [0, 1], b, [1]); // Shape [2]
+ * ```
+ */
+declare function einsum(...args: (ArrayLike | number[])[]): Array;
+/**
+ * Compute the inner product of two arrays.
+ *
+ * Unlike `jax.numpy.matmul()` or `jax.numpy.dot()`, this always performs a
+ * contraction on the last axis.
+ *
+ * Returned array has shape `[...x.shape[:-1], ...y.shape[:-1]]`.
+ */
+declare function inner(x: ArrayLike, y: ArrayLike): Array;
+/**
+ * Compute the outer product of two arrays.
+ *
+ * If the input arrays are not 1D, they will be flattened. Returned array will
+ * be of shape `[x.size, y.size]`.
+ */
+declare function outer(x: ArrayLike, y: ArrayLike): Array;
+/** Vector dot product of two arrays along a given axis. */
+declare function vecdot(x: ArrayLike, y: ArrayLike, {
+  axis
+}?: {
+  axis?: number;
+}): Array;
+/**
+ * Return the dot product of two vectors.
+ *
+ * Like vecdot() but flattens the arguments first into vectors.
+ */
+declare function vdot(x: ArrayLike, y: ArrayLike): Array;
+/**
+ * Return a tuple of coordinate matrices from coordinate vectors.
+ *
+ * Make N-D coordinate arrays for vectorized evaluations of N-D scalar/vector
+ * fields over N-D grids, given one-dimensional coordinate arrays x1, x2,…, xn.
+ */
+declare function meshgrid(xs: Array[], {
+  indexing
+}?: {
+  indexing?: "xy" | "ij";
+}): Array[];
+/**
+ * Return an array with ones on and below the diagonal and zeros elsewhere.
+ *
+ * If `k` is provided, it specifies the sub-diagonal on and below which the
+ * array is filled with ones. `k=0` is the main diagonal, `k<0` is below it, and
+ * `k>0` is above it.
+ */
+declare function tri(n: number, m?: number, k?: number, {
+  dtype,
+  device
+}?: DTypeAndDevice): Array;
+/** Return the lower triangle of an array. Must be of dimension >= 2. */
+declare function tril(a: ArrayLike, k?: number): Array;
+/** Return the upper triangle of an array. Must be of dimension >= 2. */
+declare function triu(a: ArrayLike, k?: number): Array;
+/**
+ * Clip (limit) the values in an array.
+ *
+ * Given an interval, values outside the interval are clipped to the interval
+ * edges. For example, if an interval of [0, 1] is specified, values smaller
+ * than 0 become 0, and values larger than 1 become 1.
+ *
+ * If either bound is undefined, it is ignored.
+ */
+declare function clip(a: ArrayLike, min?: ArrayLike, max?: ArrayLike): Array;
+/**
+ * Calculate the absolute value element-wise.
+ *
+ * This is the same function as `jax.numpy.abs()`.
+ */
+declare function absolute(x: ArrayLike): Array;
+/** @function Alias of `jax.numpy.absolute()`. */
+declare const abs: typeof absolute;
+/** 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`.
+ */
+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). */
+declare function tan(x: ArrayLike): Array;
+/** Element-wise inverse cosine function (inverse of cos). */
+declare function acos(x: ArrayLike): Array;
+/**
+ * @function
+ * 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.
+ */
+declare const hypot: OwnedFunction<(x1: ArrayLike, x2: ArrayLike) => Array>;
+/**
+ * @function
+ * Element-wise arc tangent of y/x with correct quadrant.
+ *
+ * Returns the angle in radians between the positive x-axis and the point (x, y).
+ * The result is in the range [-π, π].
+ *
+ * Uses numerically stable formulas:
+ * - When x >= 0: atan2(y, x) = 2 * atan(y / (sqrt(x^2 + y^2) + x))
+ * - When x < 0:  atan2(y, x) = 2 * atan((sqrt(x^2 + y^2) - x) / y)
+ *
+ * The output is ill-defined when both x and y are zero.
+ */
+declare const atan2: OwnedFunction<(y: ArrayLike, x: ArrayLike) => Array>;
+/** @function Alias of `jax.numpy.acos()`. */
+declare const arccos: typeof acos;
+/** @function Alias of `jax.numpy.atan()`. */
+declare const arctan: (x: ArrayLike) => Array;
+/** @function Alias of `jax.numpy.atan2()`. */
+declare const arctan2: OwnedFunction<(y: ArrayLike, x: ArrayLike) => Array>;
+/** Element-wise subtraction, with broadcasting. */
+declare function subtract(x: ArrayLike, y: ArrayLike): Array;
+/** Calculates the floating-point division of x by y element-wise. */
+declare function trueDivide(x: ArrayLike, y: ArrayLike): Array;
+/**
+ * @function
+ * Calculate element-wise floating-point modulo operation.
+ */
+declare const fmod: OwnedFunction<(x: ArrayLike, y: ArrayLike) => Array>;
+/**
+ * @function
+ * Calculate element-wise remainder of the division (matches sign of y).
+ */
+declare const remainder: OwnedFunction<(x: ArrayLike, y: ArrayLike) => Array>;
+/** @function Alias of `jax.numpy.trueDivide()`. */
+declare const divide: typeof trueDivide;
+/** Round input to the nearest integer towards zero. */
+declare function trunc(x: ArrayLike): Array;
+/**
+ * Compute `x1 * 2 ** x2` as a standard multiplication and exponentiation.
+ *
+ * This is the inverse of `frexp()`.
+ */
+declare function ldexp(x1: ArrayLike, x2: ArrayLike): Array;
+/**
+ * Decompose floating-point values into mantissa and two's exponent.
+ *
+ * The mantissa is returned in the range `(-1, 1)` with magnitude `>= 0.5` if
+ * `x != 0`, and the exponent is an integer such that
+ * `x = mantissa * 2**exponent`.
+ */
+declare function frexp(x: ArrayLike): [Array, Array];
+/** Calculate `2**p` for all p in the input array. */
+declare function exp2(p: ArrayLike): Array;
+/** Return the base-2 logarithm of x, element-wise. */
+declare function log2(x: ArrayLike): Array;
+/** Return the base-10 logarithm of x, element-wise. */
+declare function log10(x: ArrayLike): Array;
+/** Calculate `exp(x) - 1` element-wise. */
+declare function expm1(x: ArrayLike): Array;
+/** Calculate the natural logarithm of `1 + x` element-wise. */
+declare function log1p(x: ArrayLike): Array;
+/** Convert angles from degrees to radians. */
+declare function deg2rad(x: ArrayLike): Array;
+/** @function Alias of `jax.numpy.deg2rad()`. */
+declare const radians: typeof deg2rad;
+/** Convert angles from radians to degrees. */
+declare function rad2deg(x: ArrayLike): Array;
+/** @function Alias of `jax.numpy.rad2deg()`. */
+declare const degrees: typeof rad2deg;
+/**
+ * @function
+ * Computes first array raised to power of second array, element-wise.
  */
-declare function tile(a: ArrayLike, reps: number | number[]): Array;
+declare const power: OwnedFunction<(x1: ArrayLike, x2: ArrayLike) => Array>;
+/** @function Alias of `jax.numpy.power()`. */
+declare const pow: OwnedFunction<(x1: ArrayLike, x2: ArrayLike) => Array>;
+/** @function Calculate the element-wise cube root of the input array. */
+declare const cbrt: OwnedFunction<(x: ArrayLike) => Array>;
 /**
- * Broadcast an array to a shape, with NumPy-style broadcasing rules.
+ * @function
+ * Calculate element-wise hyperbolic sine of input.
  *
- * In other words, this lets you append axes to the left, and/or expand
- * dimensions where the shape is 1.
+ * `sinh(x) = (exp(x) - exp(-x)) / 2`
  */
-declare function broadcastTo(a: ArrayLike, shape: number[]): Array;
-/** Broadcast input shapes to a common output shape. */
-declare function broadcastShapes(...shapes: number[][]): number[];
-/** Broadcast arrays to a common shape. */
-declare function broadcastArrays(...arrays: ArrayLike[]): Array[];
+declare const sinh: OwnedFunction<(x: ArrayLike) => Array>;
 /**
- * Return specified diagonals.
- *
- * If a is 2D, return the diagonal of the array with the given offset. If a is
- * 3D or higher, compute diagonals along the two given axes (default: 0, 1).
+ * @function
+ * Calculate element-wise hyperbolic cosine of input.
  *
- * This returns a view over the existing array. The shape of the resulting array
- * is determined by removing the two axes along which the diagonal is taken,
- * then appending a new axis to the right with holding the diagonals.
+ * `cosh(x) = (exp(x) + exp(-x)) / 2`
  */
-declare function diagonal(a: ArrayLike, offset?: number, axis1?: number, axis2?: number): Array;
+declare const cosh: OwnedFunction<(x: ArrayLike) => Array>;
 /**
- * Extract a diagonal or construct a diagonal array.
+ * @function
+ * Calculate element-wise hyperbolic tangent of input.
  *
- * If v is a 2D array, return the k-th diagonal of v (as a view). If v is a 1D
- * array, return a 2D array with v on the k-th diagonal.
+ * `tanh(x) = sinh(x)/cosh(x) = (exp(x) - exp(-x)) / (exp(x) + exp(-x))`
  */
-declare function diag(v: ArrayLike, k?: number): Array;
-/** Return if two arrays are element-wise equal within a tolerance. */
-declare function allclose(actual: Parameters<typeof array>[0], expected: Parameters<typeof array>[0], options?: {
-  rtol?: number;
-  atol?: number;
-}): boolean;
-/** Matrix product of two arrays. */
-declare function matmul(x: ArrayLike, y: ArrayLike): Array;
-/** Dot product of two arrays. */
-declare function dot(x: ArrayLike, y: ArrayLike): Array;
+declare const tanh: OwnedFunction<(x: ArrayLike) => Array>;
 /**
- * Compute the inner product of two arrays.
- *
- * Unlike `jax.numpy.matmul()` or `jax.numpy.dot()`, this always performs a
- * contraction on the last axis.
+ * @function
+ * Calculate element-wise inverse hyperbolic sine of input.
  *
- * Returned array has shape `[...x.shape[:-1], ...y.shape[:-1]]`.
+ * `arcsinh(x) = ln(x + sqrt(x^2 + 1))`
  */
-declare function inner(x: ArrayLike, y: ArrayLike): Array;
+declare const arcsinh: OwnedFunction<(x: ArrayLike) => Array>;
 /**
- * Compute the outer product of two arrays.
+ * @function
+ * Calculate element-wise inverse hyperbolic cosine of input.
  *
- * If the input arrays are not 1D, they will be flattened. Returned array will
- * be of shape `[x.size, y.size]`.
+ * `arccosh(x) = ln(x + sqrt(x^2 - 1))`
  */
-declare function outer(x: ArrayLike, y: ArrayLike): Array;
-/** Vector dot product of two arrays along a given axis. */
-declare function vecdot(x: ArrayLike, y: ArrayLike, {
-  axis
-}?: {
-  axis?: number;
-}): Array;
+declare const arccosh: OwnedFunction<(x: ArrayLike) => Array>;
 /**
- * Return the dot product of two vectors.
+ * @function
+ * Calculate element-wise inverse hyperbolic tangent of input.
  *
- * Like vecdot() but flattens the arguments first into vectors.
+ * `arctanh(x) = 0.5 * ln((1 + x) / (1 - x))`
  */
-declare function vdot(x: ArrayLike, y: ArrayLike): Array;
+declare const arctanh: OwnedFunction<(x: ArrayLike) => Array>;
+/** @function Alias of `jax.numpy.arcsinh()`. */
+declare const asinh: OwnedFunction<(x: ArrayLike) => Array>;
+/** @function Alias of `jax.numpy.arccosh()`. */
+declare const acosh: OwnedFunction<(x: ArrayLike) => Array>;
+/** @function Alias of `jax.numpy.arctanh()`. */
+declare const atanh: OwnedFunction<(x: ArrayLike) => Array>;
 /**
- * Return a tuple of coordinate matrices from coordinate vectors.
+ * Compute the variance of an array.
  *
- * Make N-D coordinate arrays for vectorized evaluations of N-D scalar/vector
- * fields over N-D grids, given one-dimensional coordinate arrays x1, x2,…, xn.
- */
-declare function meshgrid(xs: Array[], {
-  indexing
-}?: {
-  indexing?: "xy" | "ij";
-}): Array[];
-/**
- * Return an array with ones on and below the diagonal and zeros elsewhere.
+ * The variance is computed for the flattened array by default, otherwise over
+ * the specified axis.
  *
- * If `k` is provided, it specifies the sub-diagonal on and below which the
- * array is filled with ones. `k=0` is the main diagonal, `k<0` is below it, and
- * `k>0` is above it.
+ * If `correction` is provided, the divisor in calculation is `N - correction`,
+ * where `N` represents the number of elements (e.g., for Bessel's correction).
  */
-declare function tri(n: number, m?: number, k?: number, {
-  dtype,
-  device
-}?: DTypeAndDevice): Array;
-/** Return the lower triangle of an array. Must be of dimension >= 2. */
-declare function tril(a: ArrayLike, k?: number): Array;
-/** Return the upper triangle of an array. Must be of dimension >= 2. */
-declare function triu(a: ArrayLike, k?: number): Array;
+declare function var_(x: ArrayLike, axis?: Axis, opts?: {
+  mean?: ArrayLike;
+  correction?: number;
+} & ReduceOpts): Array;
 /**
- * Clip (limit) the values in an array.
+ * Compute the standard deviation of an array.
  *
- * Given an interval, values outside the interval are clipped to the interval
- * edges. For example, if an interval of [0, 1] is specified, values smaller
- * than 0 become 0, and values larger than 1 become 1.
+ * The standard deviation is computed for the flattened array by default,
+ * otherwise over the specified axis.
  *
- * If either bound is undefined, it is ignored.
+ * If `correction` is provided, the divisor in calculation is `N - correction`,
+ * where `N` represents the number of elements (e.g., for Bessel's correction).
  */
-declare function clip(a: ArrayLike, min?: ArrayLike, max?: ArrayLike): Array;
+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;
 /**
- * Calculate the absolute value element-wise.
+ * @function
+ * Test element-wise for finite values (not infinity or NaN).
+ */
+declare const isfinite: OwnedFunction<(x: ArrayLike) => Array>;
+//# sourceMappingURL=numpy.d.ts.map
+declare namespace tree_d_exports {
+  export { JsTree, JsTreeDef, MapJsTree, NodeType, dispose, flatten, leaves, map, ref, structure, unflatten };
+}
+declare enum NodeType {
+  Array = "Array",
+  Object = "Object",
+  Leaf = "Leaf",
+}
+/** Analog to the JAX "pytree" object, but for JavaScript. */
+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 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. */
+declare class JsTreeDef {
+  readonly nodeType: NodeType;
+  readonly nodeMetadata: any;
+  readonly childTreedefs: JsTreeDef[];
+  static leaf: 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. */
+  equals(other: JsTreeDef): boolean;
+}
+/** Flatten a structured object, returning the tree definition. */
+declare function flatten<T>(tree: JsTree<T>): [T[], JsTreeDef];
+/** Get the leaves of a tree. */
+declare function leaves<T>(tree: JsTree<T>): T[];
+/** Get the treedef for a tree. */
+declare function structure<T>(tree: JsTree<T>): JsTreeDef;
+/** Reconstruct a structured object from the flattened representation. */
+declare function unflatten<T>(treedef: JsTreeDef, leaves: Iterable<T>): JsTree<T>;
+/** Maps a multi-input function over pytree args to produce a new pytree. */
+declare function map<T, U, Tree extends JsTree<T>>(fn: (...args: T[]) => U, tree: Tree, ...rest: Tree[]): MapJsTree<Tree, T, U>;
+/** Take a reference of every array in a tree. */
+declare function ref<Tree extends JsTree<Array>>(tree: Tree): Tree;
+/** Dispose every array in a tree. */
+declare function dispose<Tree extends JsTree<Array>>(tree: Tree | null | undefined): void;
+//#endregion
+//#region src/frontend/convolution.d.ts
+/** Definition of a general dilated convolution. Should be valid on creation. */
+interface ConvParams {
+  vmapDims: number;
+  strides: number[];
+  padding: [number, number][];
+  lhsDilation: number[];
+  rhsDilation: number[];
+}
+/**
+ * Check that the shapes and parameters passed to convolution are valid.
+ * Expected shapes of the lhs and rhs of the convolution are:
  *
- * This is the same function as `jax.numpy.abs()`.
+ * - `lhsShape = [*vmapDims, batchSize, inChannels, spatialDims...]`
+ * - `rhsShape = [*vmapDims, outChannels, inChannels, kernelSize...]`
+ *
+ * If the check succeeds, returns the output shape.
  */
-declare function absolute(x: ArrayLike): Array;
-/** @function Alias of `jax.numpy.absolute()`. */
-declare const abs: typeof absolute;
-/** Return an element-wise indication of sign of the input. */
-declare function sign(x: ArrayLike): Array;
+//#endregion
+//#region src/frontend/jaxpr.d.ts
 /**
- * Return the Hamming window of size M, a taper with a weighted cosine bell.
+ * Function callback with an associated dispose() method.
  *
- * `w(n) = 0.54 - 0.46 * cos(2πn/(M-1))` for `0 <= n <= M-1`.
+ * The dispose() method should be called to clean up any tracer resources needed
+ * by the function after the last time it is called.
  */
-declare function hamming(M: number): Array;
+type OwnedFunction<F extends Function> = F & {
+  dispose: () => void;
+};
+/** Variable in a Jaxpr expression. */
+declare class Var {
+  #private;
+  readonly id: number;
+  readonly aval: ShapedArray;
+  constructor(aval: ShapedArray);
+  toString(): string;
+}
+/** Literal in a Jaxpr expression. Currently, only scalars are supported. */
+declare class Lit {
+  readonly value: number;
+  readonly aval: ShapedArray;
+  get dtype(): DType;
+  constructor(aval: AbstractValue, value: number);
+}
+type Atom = Var | Lit;
+declare class VarPrinter {
+  #private;
+  names: Map<Var, string>;
+  name(v: Var): string;
+  nameType(v: Var): string;
+}
+/** A single statement / binding in a Jaxpr, in ANF form. */
+declare class JaxprEqn {
+  readonly primitive: Primitive;
+  readonly inputs: Atom[];
+  readonly params: Record<string, any>;
+  readonly outBinders: Var[];
+  constructor(primitive: Primitive, inputs: Atom[], params: Record<string, any>, outBinders: Var[]);
+  pprint(usedVars?: Set<Var>, vp?: VarPrinter): PPrint;
+  toString(): string;
+}
+/** Typed intermediate representation for traced computations. */
+declare class Jaxpr implements FpHashable {
+  #private;
+  readonly inBinders: Var[];
+  readonly eqns: JaxprEqn[];
+  readonly outs: Atom[];
+  constructor(inBinders: Var[], eqns: JaxprEqn[], outs: Atom[]);
+  pprint(): PPrint;
+  toString(): string;
+  /**
+   * Gets a hash of this Jaxpr.
+   *
+   * Var identity is not considered in the hash, so two Jaxprs with the same
+   * order of assignments and operators but different variable IDs will resolve
+   * to the same hash (and toString representation).
+   */
+  getHash(): bigint;
+  hash(state: FpHash): void;
+  /**
+   * Produce a simplified Jaxpr with basic optimizations applied.
+   *  - Trim away unused variables.
+   *  - Fold away *1, *0, or +0 operations against literals.
+   *  - Remove no-op movement operations.
+   */
+  simplify(): Jaxpr;
+  /** Flattens nested JitCall in a Jaxpr. Useful for handling jit-of-jit. */
+  flatten(): Jaxpr;
+}
+/** @inline */
+type JitOpts = {
+  staticArgnums?: number[];
+};
+//#endregion
+//#region src/frontend/core.d.ts
 /**
- * Return the Hann window of size M, a taper with a weighted cosine bell.
+ * Frontend primitive operations, which are lowered into Kernel objects before
+ * being dispatched to the backend.
  *
- * `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). */
-declare function tan(x: ArrayLike): Array;
-/** Element-wise inverse cosine function (inverse of cos). */
-declare function acos(x: ArrayLike): Array;
-/**
- * @function
- * Return element-wise hypotenuse for the given legs of a right triangle.
+ * Any operation between arrays can be described in these parts. This is also
+ * the set of primitives that can occur in Jaxpr programs, and the level at
+ * which transformations like vmap, grad, and jvp occur. They are loosely based
+ * on [XLA](https://openxla.org/xla/operation_semantics).
  *
- * 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.
+ * All n-ary operations support broadcasting, with NumPy semantics.
  */
-declare const hypot: OwnedFunction<(x1: ArrayLike, x2: ArrayLike) => Array>;
+declare enum Primitive {
+  Add = "add",
+  Mul = "mul",
+  Idiv = "idiv",
+  Mod = "mod",
+  // uses sign of dividend, C-style, matches JS but not Python
+  Neg = "neg",
+  Reciprocal = "reciprocal",
+  Floor = "floor",
+  Ceil = "ceil",
+  StopGradient = "stop_gradient",
+  Cast = "cast",
+  Bitcast = "bitcast",
+  RandomBits = "random_bits",
+  Sin = "sin",
+  Cos = "cos",
+  Asin = "asin",
+  Atan = "atan",
+  Exp = "exp",
+  Log = "log",
+  Erf = "erf",
+  Erfc = "erfc",
+  Sqrt = "sqrt",
+  Min = "min",
+  Max = "max",
+  Reduce = "reduce",
+  Dot = "dot",
+  // sum(x*y, axis=-1)
+  Conv = "conv",
+  // see lax.conv_general_dilated
+  Pool = "pool",
+  PoolTranspose = "pool_transpose",
+  Compare = "compare",
+  Where = "where",
+  Transpose = "transpose",
+  Broadcast = "broadcast",
+  Reshape = "reshape",
+  Flip = "flip",
+  Shrink = "shrink",
+  Pad = "pad",
+  Gather = "gather",
+  JitCall = "jit_call",
+}
+interface PrimitiveParamsImpl extends Record<Primitive, Record<string, any>> {
+  [Primitive.Cast]: {
+    dtype: DType;
+  };
+  [Primitive.Bitcast]: {
+    dtype: DType;
+  };
+  [Primitive.Reduce]: {
+    op: AluOp;
+    axis: number[];
+  };
+  [Primitive.Conv]: ConvParams;
+  [Primitive.Pool]: {
+    window: number[];
+    strides: number[];
+  };
+  [Primitive.PoolTranspose]: {
+    inShape: number[];
+    window: number[];
+    strides: number[];
+  };
+  [Primitive.Compare]: {
+    op: CompareOp;
+  };
+  [Primitive.Transpose]: {
+    perm: number[];
+  };
+  [Primitive.Broadcast]: {
+    shape: number[];
+    axis: number[];
+  };
+  [Primitive.RandomBits]: {
+    shape: number[];
+    mode: "xor" | 0 | 1;
+  };
+  [Primitive.Reshape]: {
+    shape: number[];
+  };
+  [Primitive.Flip]: {
+    axis: number[];
+  };
+  [Primitive.Shrink]: {
+    slice: Pair[];
+  };
+  [Primitive.Pad]: {
+    width: Pair[];
+  };
+  [Primitive.Gather]: {
+    axis: number[];
+    outDim: number;
+  };
+  [Primitive.JitCall]: {
+    name: string;
+    jaxpr: Jaxpr;
+    numConsts: number;
+  };
+}
+/** Type of parameters taken by each primitive. */
+type PrimitiveParams<T extends Primitive> = T extends keyof PrimitiveParamsImpl ? PrimitiveParamsImpl[T] : Record<string, never>;
+declare enum CompareOp {
+  Less = "less",
+  Equal = "equal",
+  NotEqual = "not_equal",
+  LessEqual = "less_equal",
+}
+/** @inline */
+type Axis = number | number[] | null;
+/** @inline */
+type ReduceOpts = {
+  keepdims?: boolean;
+};
+type MainTrace = {
+  level: number;
+  traceType: new (main: MainTrace) => Trace;
+  globalData: any | null;
+};
 /**
- * @function
- * Element-wise arc tangent of y/x with correct quadrant.
- *
- * Returns the angle in radians between the positive x-axis and the point (x, y).
- * The result is in the range [-π, π].
- *
- * Uses numerically stable formulas:
- * - When x >= 0: atan2(y, x) = 2 * atan(y / (sqrt(x^2 + y^2) + x))
- * - When x < 0:  atan2(y, x) = 2 * atan((sqrt(x^2 + y^2) - x) / y)
- *
- * The output is ill-defined when both x and y are zero.
+ * Push an interpreter onto the trace stack. Use this like:
+ * `using main = newMain(...);`
  */
-declare const atan2: OwnedFunction<(y: ArrayLike, x: ArrayLike) => Array>;
-/** @function Alias of `jax.numpy.acos()`. */
-declare const arccos: typeof acos;
-/** @function Alias of `jax.numpy.atan()`. */
-declare const arctan: (x: ArrayLike) => Array;
-/** @function Alias of `jax.numpy.atan2()`. */
-declare const arctan2: OwnedFunction<(y: ArrayLike, x: ArrayLike) => Array>;
-/** Element-wise subtraction, with broadcasting. */
-declare function subtract(x: ArrayLike, y: ArrayLike): Array;
-/** Calculates the floating-point division of x by y element-wise. */
-declare function trueDivide(x: ArrayLike, y: ArrayLike): Array;
-/** @function Alias of `jax.numpy.trueDivide()`. */
-declare const divide: typeof trueDivide;
-/** Round input to the nearest integer towards zero. */
-declare function trunc(x: ArrayLike): Array;
-/** Calculate `2**p` for all p in the input array. */
-declare function exp2(p: ArrayLike): Array;
-/** Return the base-2 logarithm of x, element-wise. */
-declare function log2(x: ArrayLike): Array;
-/** Return the base-10 logarithm of x, element-wise. */
-declare function log10(x: ArrayLike): Array;
-/** Calculate `exp(x) - 1` element-wise. */
-declare function expm1(x: ArrayLike): Array;
-/** Calculate the natural logarithm of `1 + x` element-wise. */
-declare function log1p(x: ArrayLike): Array;
-/** Convert angles from degrees to radians. */
-declare function deg2rad(x: ArrayLike): Array;
-/** @function Alias of `jax.numpy.deg2rad()`. */
-declare const radians: typeof deg2rad;
-/** Convert angles from radians to degrees. */
-declare function rad2deg(x: ArrayLike): Array;
-/** @function Alias of `jax.numpy.rad2deg()`. */
-declare const degrees: typeof rad2deg;
+
+type TracerValue = Tracer | number | boolean;
+declare abstract class Trace {
+  readonly main: MainTrace;
+  constructor(main: MainTrace);
+  abstract pure(val: TracerValue): Tracer;
+  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;
+}
 /**
- * @function
- * Computes first array raised to power of second array, element-wise.
+ * 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 const power: OwnedFunction<(x1: ArrayLike, x2: ArrayLike) => Array>;
-/** @function Alias of `jax.numpy.power()`. */
-declare const pow: OwnedFunction<(x1: ArrayLike, x2: ArrayLike) => Array>;
-/** @function Calculate the element-wise cube root of the input array. */
-declare const cbrt: OwnedFunction<(x: ArrayLike) => Array>;
+
+declare abstract class Tracer {
+  /** @ignore */
+  readonly _trace: Trace;
+  constructor(trace: Trace);
+  abstract get aval(): AbstractValue;
+  abstract toString(): string;
+  /**
+   * Access an array by reference, incrementing the reference count.
+   *
+   * jax-js handles freeing arrays by using "move" semantics, like in Rust/C++.
+   * Whenever you pass an array into a function, that function should consume
+   * the array, and it will no longer be usable. For example, if you had:
+   *
+   * ```
+   * const x = np.array([1, 2, 3]);
+   * const y = np.add(x, x);
+   * ```
+   *
+   * The second line does not work because the first parameter consumes `x`, and
+   * then the second parameter will already have been freed / disposed.
+   *
+   * To fix this, you can write:
+   *
+   * ```
+   * const y = np.add(x.ref, x);
+   * ```
+   *
+   * Under the hood, every access to `.ref` increments the internal reference
+   * count of the array. The reference count starts at 1. When it hits 0, the
+   * memory behind the array is freed.
+   */
+  abstract get ref(): this;
+  /**
+   * Manually decrement the reference count of the array.
+   *
+   * Arrays are created with reference count 1. Whenever it is used as argument
+   * to a function or other operation, it is disposed (i.e., reference count
+   * decreases by 1) automatically. Whenever a `.ref` is created, the reference
+   * count increases.
+   *
+   * You generally don't need to call this function directly since arrays are
+   * automatically disposed after being passed into an operation. One common
+   * exception is when writing a function and ignoring one of its arguments. In
+   * that case, by convention you should dispose of that argument manually.
+   *
+   * ```
+   * function myCustomOperation(a: np.Array, b: np.Array) {
+   *   b.dispose(); // Needed to satisfy "move" rules.
+   *   return a.add(1);
+   * }
+   * ```
+   */
+  abstract dispose(): void;
+  /** The shape of the array. */
+  get shape(): number[];
+  /** The total number of elements in the array. */
+  get size(): number;
+  /** 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 */
+  fullLower(): Tracer;
+  neg(): this;
+  add(other: this | TracerValue): this;
+  mul(other: this | TracerValue): this;
+  greater(other: this | TracerValue): this;
+  less(other: this | TracerValue): this;
+  equal(other: this | TracerValue): this;
+  notEqual(other: this | TracerValue): this;
+  greaterEqual(other: this | TracerValue): this;
+  lessEqual(other: this | TracerValue): this;
+  /** Sum of the elements of the array over a given axis, or axes. */
+  sum(axis?: Axis, opts?: ReduceOpts): this;
+  /** Product of the array elements over a given axis. */
+  prod(axis?: Axis, opts?: ReduceOpts): this;
+  /** Compute the average of the array elements along the specified axis. */
+  mean(axis?: Axis, opts?: ReduceOpts): this;
+  /** Permute the dimensions of an array. Defaults to reversing the axis order. */
+  transpose(perm?: number[]): this;
+  /**
+   * Give a new shape to an array without changing its data.
+   *
+   * One shape dimension can be -1. In this case, the value is inferred from the
+   * length of the array and remaining dimensions.
+   */
+  reshape(shape: number | number[]): this;
+  /** Copy the array and cast to a specified dtype. */
+  astype(dtype: DType): this;
+  /** Subtract an array from this one. */
+  sub(other: this | TracerValue): this;
+  /** Divide an array by this one. */
+  div(other: this | TracerValue): this;
+  /** Return specified diagonals. See `numpy.diagonal` for full docs. */
+  diagonal(offset?: number, axis1?: number, axis2?: number): this;
+  /** Flatten the array without changing its data. */
+  flatten(): this;
+  /** Flatten the array without changing its data. */
+  ravel(): this;
+  /**
+   * Iterate over the first dimension of this array, returning slices.
+   *
+   * This can be used to destructure arrays. For example:
+   *
+   * ```js
+   * let x = np.array([[1, 2], [3, 4]]);
+   * let [a, b] = x;
+   * console.log(a.js()); // [1, 2]
+   * console.log(b.js()); // [3, 4]
+   * ```
+   */
+  [Symbol.iterator](): IterableIterator<this>;
+  /**
+   * Slice an array along one or more axes.
+   *
+   * This is the equivalent of slicing in Python, e.g. `x[1:3, 2, :, None]`. To
+   * mimic this in JavaScript, we would write:
+   *
+   * ```js
+   * x.slice([1, 3], 2, [], null);
+   * ```
+   *
+   * The `slice` method accepts a variable number of arguments, each of which
+   * can be a number, an empty array, a single-element array, a two-element
+   * array, or `null`. The arguments are interpreted as follows:
+   *
+   * - A number `n` means to access the `n`-th element along that axis, removing
+   *   that axis from the resulting shape.
+   * - An empty array `[]` means to keep that axis as-is, like `:` in Python.
+   * - A single-element array `[i]` means to start slicing from index `i`
+   *   (inclusive) to the end of the axis, like `x[i:]`.
+   * - A two-element array `[i, j]` means to slice from index `i` (inclusive)
+   *   to index `j` (exclusive), like `x[i:j]`.
+   * - `null` means to add a new axis at that position, like `np.newaxis`.
+   *
+   * Like in Python, negative indices are supported, which count from the end of
+   * the axis. For example, `-1` means the last element.
+   *
+   * Strided slices are not yet implemented, so you cannot write `x[::2]` or
+   * similar.
+   *
+   * Advanced indexing by integer arrays is also supported. This translates to
+   * the "gather" primitive, and it allows you to access specific elements of
+   * the array by integer indices stored in another array.
+   */
+  slice(...index: (number | [] | [number] | Pair | null | Tracer)[]): this;
+}
+declare class ShapedArray implements AbstractValue {
+  readonly shape: number[];
+  readonly dtype: DType;
+  readonly weakType: boolean;
+  constructor(shape: number[], dtype: DType, weakType: boolean);
+  static fromAval(aval: AbstractValue): ShapedArray;
+  get ndim(): number;
+  toString(): string;
+  equals(other: ShapedArray): boolean;
+}
+//#endregion
+//#region src/frontend/array.d.ts
+type ArrayLike = Array | number | boolean;
+/** Version of pureArray with fudged types. */
+
 /**
- * @function
- * Calculate element-wise hyperbolic sine of input.
+ * An executable operation that will be dispatched to the backend.
  *
- * `sinh(x) = (exp(x) - exp(-x)) / 2`
+ * This holds a reference to all input buffers used in the operation. After the
+ * operation is dispatched, the references should be released.
  */
-declare const sinh: OwnedFunction<(x: ArrayLike) => Array>;
+declare class PendingExecute {
+  #private;
+  readonly backend: Backend;
+  readonly kernel: Kernel;
+  readonly inputs: Slot[];
+  readonly outputs: Slot[];
+  prepared: Executable | null;
+  submitted: boolean;
+  constructor(backend: Backend, kernel: Kernel, inputs: Slot[], outputs: Slot[]);
+  updateRc(delta: number): void;
+  prepare(): Promise<void>;
+  prepareSync(): void;
+  submit(): void;
+}
+/** @inline */
+type DTypeAndDevice = {
+  dtype?: DType;
+  device?: Device;
+};
+type ArrayConstructorArgs = {
+  source: AluExp | Slot;
+  st: ShapeTracker;
+  dtype: DType;
+  weakType: boolean;
+  backend: Backend;
+  committed: boolean;
+  pending?: Iterable<PendingExecute>;
+};
 /**
- * @function
- * Calculate element-wise hyperbolic cosine of input.
+ * A multidimensional numeric array with data stored on CPU or GPU.
  *
- * `cosh(x) = (exp(x) + exp(-x)) / 2`
+ * 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
+ * this into your code's namespace if you're already using the JavaScript
+ * "Array" type by name.
  */
-declare const cosh: OwnedFunction<(x: ArrayLike) => Array>;
+declare class Array extends Tracer {
+  #private;
+  id: number;
+  /**
+   * @ignore
+   * Constructs an array from source, shape and backend. Note that if the source
+   * is a backend `Slot`, this constructor _takes ownership_ of the slot. It
+   * will be freed when the array is disposed.
+   */
+  constructor(args: ArrayConstructorArgs);
+  /** @ignore */
+  get aval(): ShapedArray;
+  /** Return a simple string representation of the array's dimensions. */
+  toString(): string;
+  get device(): Device;
+  get ref(): this;
+  dispose(): void;
+  /**
+   * Convert this array into a primitive value.
+   *
+   * This only works for scalars (0-dimensional arrays). It lets you get values
+   * "out" of the JAX system. For instance, if `x = np.array(5)`, then you can
+   * evaluate `x + 1` and `x ** 2` to get `6` and `25`, respectively.
+   *
+   * This method is also called for `==` equality.
+   */
+  [Symbol.toPrimitive](): any;
+  /** Realize the array and return it as data. */
+  data(): Promise<DataArray>;
+  /**
+   * Wait for this array to finish evaluation.
+   *
+   * Operations and data loading in jax-js are lazy, so this function ensures
+   * that pending operations are dispatched and fully executed before it
+   * returns.
+   *
+   * If you are mapping from `data()` or `dataSync()`, it will also trigger
+   * dispatch of operations as well.
+   *
+   * **Note:** `jax.blockUntilReady()` is a higher-level API, it calls this
+   * asynchronously for multiple arrays.
+   */
+  blockUntilReady(): Promise<Array>;
+  /**
+   * Realize the array and return it as data. This is a sync variant and not
+   * recommended for performance reasons, as it will block rendering.
+   */
+  dataSync(): DataArray;
+  /**
+   * Convert this array into a JavaScript object.
+   *
+   * This is a blocking operation that will compile all of the shaders and wait
+   * for execution to complete, synchronously. No other JavaScript code on the
+   * site will be run during shader execution.
+   *
+   * To avoid blocking, prefer `jsAsync()` when possible.
+   */
+  js(): any;
+  /** Convert this array into a JavaScript object, asynchronously. */
+  jsAsync(): Promise<any>;
+  /**
+   * Copy an element of an array to a numeric scalar and return it.
+   *
+   * Throws an error if the array does not have a single element. The array must
+   * either be rank-0, or all dimensions of the shape are 1.
+   */
+  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 | DataArray | RecursiveArray<number> | RecursiveArray<boolean>, {
+  shape,
+  dtype,
+  device
+}?: {
+  shape?: number[];
+} & DTypeAndDevice): Array;
+/** If x is a value, lift it into an array, otherwise leave it be. */
+
+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;
+/** Return a new array of given shape and type, filled with ones. */
+declare function ones(shape: number[], {
+  dtype,
+  device
+}?: 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,
+  device
+}?: DTypeAndDevice): Array;
 /**
- * @function
- * Calculate element-wise hyperbolic tangent of input.
+ * Create an identity matrix.
  *
- * `tanh(x) = sinh(x)/cosh(x) = (exp(x) - exp(-x)) / (exp(x) + exp(-x))`
+ * If numCols is not provided, it defaults to numRows, i.e., a square identity
+ * matrix with ones on the diagonal.
  */
-declare const tanh: OwnedFunction<(x: ArrayLike) => Array>;
+declare function eye(numRows: number, numCols?: number, {
+  dtype,
+  device
+}?: DTypeAndDevice): Array;
+/** Return the identity matrix, with ones on the main diagonal. */
+declare function identity$1(n: number, {
+  dtype,
+  device
+}?: DTypeAndDevice): Array;
 /**
- * @function
- * Calculate element-wise inverse hyperbolic sine of input.
+ * Return evenly spaced values within a given interval.
  *
- * `arcsinh(x) = ln(x + sqrt(x^2 + 1))`
- */
-declare const arcsinh: OwnedFunction<(x: ArrayLike) => Array>;
-/**
- * @function
- * Calculate element-wise inverse hyperbolic cosine of input.
+ * This can be called with a varying number of arguments, just like the range()
+ * builtin function in Python.
  *
- * `arccosh(x) = ln(x + sqrt(x^2 - 1))`
- */
-declare const arccosh: OwnedFunction<(x: ArrayLike) => Array>;
-/**
- * @function
- * Calculate element-wise inverse hyperbolic tangent of input.
+ * - `arange(stop)` is equivalent to `arange(0, stop, 1)`.
+ * - `arange(start, stop)` is equivalent to `arange(start, stop, 1)`.
+ * - `arange(start, stop, step)` creates an array starting at `start`, ending
+ *   before `stop`, with a step size of `step`.
  *
- * `arctanh(x) = 0.5 * ln((1 + x) / (1 - x))`
+ * Defaults to an integer data type. This can produce unintended results when
+ * using a non-integer step, so prefer linspace() in those cases.
  */
-declare const arctanh: OwnedFunction<(x: ArrayLike) => Array>;
-/** @function Alias of `jax.numpy.arcsinh()`. */
-declare const asinh: OwnedFunction<(x: ArrayLike) => Array>;
-/** @function Alias of `jax.numpy.arccosh()`. */
-declare const acosh: OwnedFunction<(x: ArrayLike) => Array>;
-/** @function Alias of `jax.numpy.arctanh()`. */
-declare const atanh: OwnedFunction<(x: ArrayLike) => Array>;
+declare function arange(start: number, stop?: number, step?: number, {
+  dtype,
+  device
+}?: DTypeAndDevice): Array;
 /**
- * Compute the variance of an array.
+ * Return evenly spaced numbers over a specified interval.
  *
- * The variance is computed for the flattened array by default, otherwise over
- * the specified axis.
+ * Returns _num_ evenly spaced samples, calculated over the interval
+ * [`start`, `stop`]. The endpoint `stop` is included in the result by default,
+ * but this is controlled by the `endpoint` parameter.
  *
- * If `correction` is provided, the divisor in calculation is `N - correction`,
- * where `N` represents the number of elements (e.g., for Bessel's correction).
+ * The default data type is Float32. Use arange() for integer steps.
  */
-declare function var_(x: ArrayLike, axis?: Axis, opts?: {
-  mean?: ArrayLike;
-  correction?: number;
-} & ReduceOpts): Array;
+declare function linspace(start: number, stop: number, num?: number, endpoint?: boolean, {
+  dtype,
+  device
+}?: DTypeAndDevice): Array;
+declare namespace lax_d_exports {
+  export { DotDimensionNumbers, PaddingType, conv, convGeneralDilated, convWithGeneralPadding, dot, erf, erfc, reduceWindow, stopGradient };
+}
 /**
- * Compute the standard deviation of an array.
+ * Dimension numbers for general `dot()` primitive.
  *
- * The standard deviation is computed for the flattened array by default,
- * otherwise over the specified axis.
+ * Contracting dimensions act as a tensor contraction (reduction) along the
+ * given axis. They must be the same size in both operands. Batch dimensions
+ * are treated as vectorized, leading batch dimensions.
  *
- * If `correction` is provided, the divisor in calculation is `N - correction`,
- * where `N` represents the number of elements (e.g., for Bessel's correction).
- */
-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).
+ * The return value has a shape where the first dimensions are shared batch
+ * dimensions, followed by `lhs` non-contracting dimensions, followed by
+ * `rhs` non-contracting dimensions.
  */
-declare const isfinite: OwnedFunction<(x: ArrayLike) => Array>;
-//# sourceMappingURL=numpy.d.ts.map
-//#endregion
-//#region src/frontend/jaxpr.d.ts
+type DotDimensionNumbers = {
+  lhsContractingDims?: number[];
+  rhsContractingDims?: number[];
+  lhsBatchDims?: number[];
+  rhsBatchDims?: number[];
+};
 /**
- * Function callback with an associated dispose() method.
+ * General dot product/contraction operator.
  *
- * The dispose() method should be called to clean up any tracer resources needed
- * by the function after the last time it is called.
+ * Prefer higher-level functions like `jax.numpy.dot()`, `jax.numpy.matmul()`,
+ * `jax.numpy.tensordot(), and `jax.numpy.einsum()` where possible.
  */
-type OwnedFunction<F extends Function> = F & {
-  dispose: () => void;
-};
-/** Variable in a Jaxpr expression. */
-declare class Var {
-  #private;
-  readonly id: number;
-  readonly aval: ShapedArray;
-  constructor(aval: ShapedArray);
-  toString(): string;
-}
-/** Literal in a Jaxpr expression. Currently, only scalars are supported. */
-declare class Lit {
-  readonly value: number;
-  readonly aval: ShapedArray;
-  get dtype(): DType;
-  constructor(aval: AbstractValue, value: number);
-}
-type Atom = Var | Lit;
-declare class VarPrinter {
-  #private;
-  names: Map<Var, string>;
-  name(v: Var): string;
-  nameType(v: Var): string;
-}
-/** A single statement / binding in a Jaxpr, in ANF form. */
-declare class JaxprEqn {
-  readonly primitive: Primitive;
-  readonly inputs: Atom[];
-  readonly params: Record<string, any>;
-  readonly outBinders: Var[];
-  constructor(primitive: Primitive, inputs: Atom[], params: Record<string, any>, outBinders: Var[]);
-  pprint(usedVars?: Set<Var>, vp?: VarPrinter): PPrint;
-  toString(): string;
-}
-/** Typed intermediate representation for traced computations. */
-declare class Jaxpr implements FpHashable {
-  #private;
-  readonly inBinders: Var[];
-  readonly eqns: JaxprEqn[];
-  readonly outs: Atom[];
-  constructor(inBinders: Var[], eqns: JaxprEqn[], outs: Atom[]);
-  pprint(): PPrint;
-  toString(): string;
-  /**
-   * Gets a hash of this Jaxpr.
-   *
-   * Var identity is not considered in the hash, so two Jaxprs with the same
-   * order of assignments and operators but different variable IDs will resolve
-   * to the same hash (and toString representation).
-   */
-  getHash(): bigint;
-  hash(state: FpHash): void;
-  /**
-   * Produce a simplified Jaxpr with basic optimizations applied.
-   *  - Trim away unused variables.
-   *  - Fold away *1, *0, or +0 operations against literals.
-   *  - Remove no-op movement operations.
-   */
-  simplify(): Jaxpr;
-  /** Flattens nested JitCall in a Jaxpr. Useful for handling jit-of-jit. */
-  flatten(): Jaxpr;
-}
-/** @inline */
-type JitOpts = {
-  staticArgnums?: number[];
-};
-declare namespace lax_d_exports {
-  export { PaddingType, conv, convGeneralDilated, convWithGeneralPadding, erf, erfc, reduceWindow, stopGradient };
-}
+declare function dot(lhs: Array, rhs: Array, {
+  lhsContractingDims: lc,
+  rhsContractingDims: rc,
+  lhsBatchDims: lb,
+  rhsBatchDims: rb
+}?: DotDimensionNumbers): Array;
 type PaddingType = "VALID" | "SAME" | "SAME_LOWER" | [number, number][];
 /**
  * General n-dimensional convolution operator, with optional dilation.
@@ -1636,10 +1756,12 @@ type PaddingType = "VALID" | "SAME" | "SAME_LOWER" | [number, number][];
  */
 declare function convGeneralDilated(lhs: Array, rhs: Array, windowStrides: number[], padding: PaddingType, {
   lhsDilation,
-  rhsDilation
+  rhsDilation,
+  featureGroupCount
 }?: {
   lhsDilation?: number[];
   rhsDilation?: number[];
+  featureGroupCount?: number;
 }): Array;
 /** Convenience wrapper around `convGeneralDilated`. */
 declare function convWithGeneralPadding(lhs: Array, rhs: Array, windowStrides: number[], padding: PaddingType, lhsDilation?: number[], rhsDilation?: number[]): Array;
@@ -1806,9 +1928,9 @@ declare function logSoftmax(x: ArrayLike, axis?: Axis): Array;
  *
  * Reference: https://en.wikipedia.org/wiki/LogSumExp
  */
-declare function logsumexp(x: ArrayLike, axis?: Axis): Array;
+declare function logsumexp(x: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
 /** Log-mean-exp reduction, like `jax.nn.logsumexp()` but subtracts `log(n)`. */
-declare function logmeanexp(x: ArrayLike, axis?: Axis): Array;
+declare function logmeanexp(x: ArrayLike, axis?: Axis, opts?: ReduceOpts): Array;
 /**
  * Standardizes input to zero mean and unit variance.
  *