@jax-js/jax 0.0.1 → 0.0.2

package/dist/index.d.ts

@@ -1,23 +1,24 @@
-/**
- * @file Lazy shape tracking for multidimensional tensors.
- *
- * This module provides an immutable `View` class that can be used to calculate
- * shapes of arrays as operations are applied to them, lazily.
- *
- * Some operations like reshape() may not be representable with a single view,
- * for instance, because composing reshape() with shrink() leads to a
- * non-contiguous range of memory locations. This is why `ShapeTracker` is a
- * list of views.
- *
- * Indexing into a `ShapeTracker` or `View` can be folded into shader code.
- *
- * Originally based on tinygrad's implementation of shape tracking in the
- * `tinygrad.shape` module. But this version is simplified a bit. I'm not really
- * trying to innovate on shape tracking in this library, so if I have doubts on
- * something, it'll just be copied from tinygrad (with comments).
- *
- * This file is a bit longer than the original, since Python is more concise.
- */
+//#region src/pprint.d.ts
+/** General class for pretty-printing expressions with indentation. */
+declare class PPrint {
+  readonly indents: number[];
+  readonly lines: string[];
+  constructor(indents: number[], lines: string[]);
+  /** Add a fixed amount of indentation to each line. */
+  indent(spaces: number): PPrint;
+  /** Concatenate pretty-printed expressions with newlines. */
+  concat(...items: PPrint[]): PPrint;
+  /** Stack one block to the right of another one, sharing 1 common line. */
+  stack(other: PPrint): PPrint;
+  /** Combine this block of lines into a formatted string. */
+  toString(): string;
+  static pp(s: Stringable): PPrint;
+}
+interface Stringable {
+  toString(): string;
+}
+//#endregion
+//#region src/shape.d.ts
 
 type Pair = [number, number];
 /**
@@ -30,49 +31,54 @@ type Pair = [number, number];
  *   2. Otherwise, look at this memory address: offset + ∑(strides[i] * dim[i]).
  */
 declare class View {
-    #private;
-    /** The shape of the view (size of each dimension). */
-    readonly shape: number[];
-    /** How many indices to move in buffer for each hop in one dimension. */
-    readonly strides: number[];
-    /** Offset from the start of the buffer. */
-    readonly offset: number;
-    /** Masked out subarray where data is read. All other data is zeroed. */
-    readonly mask: Pair[] | null;
-    private constructor();
-    static create(shape: number[], strides?: number[], offset?: number, mask?: Pair[] | null): View;
-    get ndim(): number;
-    get size(): number;
-    /** Whether this is a default, contiguous, unaltered view of the data (identity). */
-    get contiguous(): boolean;
-    /** Produce an AluExp for evaluating this view at an index. */
-    toAluExp(idxs: AluExp[]): [AluExp, AluExp];
-    /**
-     * Try to compose this view with another one. `this` view is applied first,
-     * followed by the argument. If this is not possible for the specific views,
-     * return `null` instead.
-     *
-     * If composable, return a combined view with the same shape as `v1`.
-     *
-     * This is very tricky. The shapes of v1 and v2 may be different, and in that
-     * case, we do some math to figure out whether they're compatible.
-     */
-    compose(v1: View): View | null;
-    /** Attempt to simplify this view into a smaller reshaped form. */
-    minify(): View;
-    /** Pad the view with zeros on each dimension. */
-    pad(arg: Pair[]): View;
-    /** Shrink the view by taking a subarray. */
-    shrink(arg: Pair[]): View;
-    /** Expand one or more axes with length "1" by repeating the data. */
-    expand(newShape: number[]): View;
-    /** Permute the axes of an array. */
-    permute(axis: number[]): View;
-    /** Flip (reverse) one or more axes of the view. */
-    flip(arg: boolean[]): View;
-    /** Reshape the view into a new shape. */
-    reshape(newShape: number[]): View | null;
+  #private;
+  /** The shape of the view (size of each dimension). */
+  readonly shape: number[];
+  /** How many indices to move in buffer for each hop in one dimension. */
+  readonly strides: number[];
+  /** Offset from the start of the buffer. */
+  readonly offset: number;
+  /** Masked out subarray where data is read. All other data is zeroed. */
+  readonly mask: Pair[] | null;
+  private constructor();
+  static create(shape: number[], strides?: number[], offset?: number, mask?: Pair[] | null): View;
+  get ndim(): number;
+  get size(): number;
+  /** Whether this is a default, contiguous, unaltered view of the data (identity). */
+  get contiguous(): boolean;
+  /** Produce an AluExp for evaluating this view at an index. */
+  toAluExp(idxs: AluExp[]): [AluExp, AluExp];
+  /**
+   * Try to compose this view with another one. `this` view is applied first,
+   * followed by the argument. If this is not possible for the specific views,
+   * return `null` instead.
+   *
+   * If composable, return a combined view with the same shape as `v1`.
+   *
+   * This is very tricky. The shapes of v1 and v2 may be different, and in that
+   * case, we do some math to figure out whether they're compatible.
+   */
+  compose(v1: View): View | null;
+  /** Attempt to simplify this view into a smaller reshaped form. */
+  minify(): View;
+  /** Pad the view with zeros on each dimension. */
+  pad(arg: Pair[]): View;
+  /** Shrink the view by taking a subarray. */
+  shrink(arg: Pair[]): View;
+  /** Expand one or more axes with length "1" by repeating the data. */
+  expand(newShape: number[]): View;
+  /** Permute the axes of an array. */
+  permute(axis: number[]): View;
+  /** Flip (reverse) one or more axes of the view. */
+  flip(arg: boolean[]): View;
+  /** Reshape the view into a new shape. */
+  reshape(newShape: number[]): View | null;
 }
+/**
+ * Find position of `offset` in each dimension within an existing shape. Like
+ * `numpy.unravel_index` in behavior.
+ */
+
 /**
  * Array shape after applying movement operations, as a series of views.
  *
@@ -80,31 +86,33 @@ declare class View {
  * shape, then used as the virtual buffer for the next view.
  */
 declare class ShapeTracker {
-    readonly views: View[];
-    constructor(views: View[]);
-    /** Compose this shape tracker with another, applying after. */
-    compose(other: ShapeTracker): ShapeTracker;
-    static fromShape(shape: number[]): ShapeTracker;
-    get contiguous(): boolean;
-    get consecutive(): boolean;
-    get lastStrides(): number[];
-    get shape(): number[];
-    get size(): number;
-    toAluExp(idxs: AluExp[]): [AluExp, AluExp];
-    simplify(): ShapeTracker;
-    pad(arg: Pair[]): ShapeTracker;
-    shrink(arg: Pair[]): ShapeTracker;
-    expand(newShape: number[]): ShapeTracker;
-    permute(axis: number[]): ShapeTracker;
-    flip(arg: boolean[]): ShapeTracker;
-    reshape(newShape: number[]): ShapeTracker;
-    /** Broadcast along the given new axes, then expand the shape. */
-    broadcast(newShape: number[], axis: number[]): ShapeTracker;
+  readonly views: View[];
+  constructor(views: View[]);
+  /** Compose this shape tracker with another, applying it after this one. */
+  compose(other: ShapeTracker): ShapeTracker;
+  static fromShape(shape: number[]): ShapeTracker;
+  get contiguous(): boolean;
+  get consecutive(): boolean;
+  get lastStrides(): number[];
+  get shape(): number[];
+  get size(): number;
+  toAluExp(idxs: AluExp[]): [AluExp, AluExp];
+  simplify(): ShapeTracker;
+  pad(arg: Pair[]): ShapeTracker;
+  shrink(arg: Pair[]): ShapeTracker;
+  expand(newShape: number[]): ShapeTracker;
+  permute(axis: number[]): ShapeTracker;
+  flip(arg: boolean[]): ShapeTracker;
+  reshape(newShape: number[]): ShapeTracker;
+  /** Broadcast along the given new axes, then expand the shape. */
+  broadcast(newShape: number[], axis: number[]): ShapeTracker;
 }
-
+//#endregion
+//#region src/utils.d.ts
+/** @inline */
 type RecursiveArray<T> = T | RecursiveArray<T>[];
 interface FpHashable {
-    hash(state: FpHash): void;
+  hash(state: FpHash): void;
 }
 /**
  * Polynomial hashes modulo p are good at avoiding collisions in expectation.
@@ -114,17 +122,20 @@ interface FpHashable {
  * See https://en.wikipedia.org/wiki/Lagrange%27s_theorem_(number_theory)
  */
 declare class FpHash {
-    #private;
-    value: bigint;
-    update(...values: (string | boolean | number | bigint | null | undefined | FpHashable)[]): this;
-    static hash(...values: (string | boolean | number | bigint | null | undefined | FpHashable)[]): bigint;
+  #private;
+  value: bigint;
+  update(...values: (string | boolean | number | bigint | null | undefined | FpHashable)[]): this;
+  static hash(...values: (string | boolean | number | bigint | null | undefined | FpHashable)[]): bigint;
 }
-
+/** Run a function while caching it inline inside a `Map`. */
+//#endregion
+//#region src/alu.d.ts
 declare enum DType {
-    Float32 = "float32",
-    Int32 = "int32",
-    Bool = "bool",
-    Complex64 = "complex64"
+  Float32 = "float32",
+  Int32 = "int32",
+  Uint32 = "uint32",
+  Bool = "bool",
+  Complex64 = "complex64",
 }
 /**
  * Mathematical expression on scalar values.
@@ -134,94 +145,105 @@ declare enum DType {
  * graph rewrite engine.
  */
 declare class AluExp implements FpHashable {
-    #private;
-    readonly op: AluOp;
-    readonly dtype: DType;
-    readonly src: AluExp[];
-    readonly arg: any;
-    constructor(op: AluOp, dtype: DType, src: AluExp[], arg?: any);
-    static add(a: AluExp, b: AluExp): AluExp;
-    static sub(a: AluExp, b: AluExp): AluExp;
-    static mul(a: AluExp, b: AluExp): AluExp;
-    static idiv(a: AluExp, b: AluExp): AluExp;
-    static mod(a: AluExp, b: AluExp): AluExp;
-    static min(a: AluExp, b: AluExp): AluExp;
-    static max(a: AluExp, b: AluExp): AluExp;
-    static sin(a: AluExp): AluExp;
-    static cos(a: AluExp): AluExp;
-    static exp(a: AluExp): AluExp;
-    static log(a: AluExp): AluExp;
-    static reciprocal(a: AluExp): AluExp;
-    static cast(dtype: DType, a: AluExp): AluExp;
-    static cmplt(a: AluExp, b: AluExp): AluExp;
-    static cmpne(a: AluExp, b: AluExp): AluExp;
-    static where(cond: AluExp, a: AluExp, b: AluExp): AluExp;
-    static const(dtype: DType, value: any): AluExp;
-    static special(dtype: DType, name: string, n: number): AluExp;
-    static variable(dtype: DType, name: string): AluExp;
-    static globalIndex(dtype: DType, gid: number, bufidx: AluExp): AluExp;
-    static globalView(dtype: DType, gid: number, st: ShapeTracker, indices: AluExp[]): AluExp;
-    static i32(value: number): AluExp;
-    static f32(value: number): AluExp;
-    static bool(value: boolean): AluExp;
-    not(): AluExp;
-    /** Compute a reasonable expression hash with low collision rate. */
-    getHash(): bigint;
-    hash(state: FpHash): void;
-    /** 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;
-    get min(): number;
-    get max(): number;
-    /**
-     * Simplify the expression by replacing any known patterns and deduping
-     * identical subexpressions.
-     */
-    simplify(cache?: Map<bigint, AluExp>): AluExp;
-    /** Resolve this to a value, or `undefined` if not possible. */
-    resolve(): any | undefined;
-    /**
-     * Evaluate the expression on CPU, returning the result.
-     *
-     * Typically you would compile the AluExp as a representation to a lower-level
-     * language. This is just to define the semantics and help debug.
-     *
-     * Note that the representation of Bool is as a number (0 or 1) here.
-     */
-    evaluate(context: Record<string, any>, globals?: (gid: number, bufidx: number) => any): any;
-    /** Get this expression in debug format as a string. */
-    toString(): string;
-    /** Generic fold() operation with a reducer over the expression tree. */
-    fold<T = void>(reducer: (exp: AluExp, mappedSrc: T[]) => T): T;
-    /** Rewrite the expression recursively using a visitor. */
-    rewrite(visitor: (exp: AluExp) => AluExp | undefined | null): AluExp;
-    /** Collect all nodes that satisfy a predicate. */
-    collect(predicate: (exp: AluExp) => boolean): AluExp[];
+  #private;
+  readonly op: AluOp;
+  readonly dtype: DType;
+  readonly src: AluExp[];
+  readonly arg: any;
+  constructor(op: AluOp, dtype: DType, src: AluExp[], arg?: any);
+  static add(a: AluExp, b: AluExp): AluExp;
+  static sub(a: AluExp, b: AluExp): AluExp;
+  static mul(a: AluExp, b: AluExp): AluExp;
+  static idiv(a: AluExp, b: AluExp): AluExp;
+  static mod(a: AluExp, b: AluExp): AluExp;
+  static min(a: AluExp, b: AluExp): AluExp;
+  static max(a: AluExp, b: AluExp): AluExp;
+  static sin(a: AluExp): AluExp;
+  static cos(a: AluExp): AluExp;
+  static exp(a: AluExp): AluExp;
+  static log(a: AluExp): AluExp;
+  static reciprocal(a: AluExp): AluExp;
+  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 cmplt(a: AluExp, b: AluExp): AluExp;
+  static cmpne(a: AluExp, b: AluExp): AluExp;
+  static where(cond: AluExp, a: AluExp, b: AluExp): AluExp;
+  static const(dtype: DType, value: any): AluExp;
+  static special(dtype: DType, name: string, n: number): AluExp;
+  static variable(dtype: DType, name: string): AluExp;
+  static globalIndex(dtype: DType, gid: number, bufidx: AluExp): AluExp;
+  static globalView(dtype: DType, gid: number, st: ShapeTracker, indices: AluExp[]): AluExp;
+  static i32(value: number): AluExp;
+  static u32(value: number): AluExp;
+  static f32(value: number): AluExp;
+  static bool(value: boolean): AluExp;
+  not(): AluExp;
+  /** Compute a reasonable expression hash with low collision rate. */
+  getHash(): bigint;
+  hash(state: FpHash): void;
+  /** 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;
+  get min(): number;
+  get max(): number;
+  /**
+   * Simplify the expression by replacing any known patterns and deduping
+   * identical subexpressions.
+   */
+  simplify(cache?: Map<bigint, AluExp>): AluExp;
+  /** Resolve this to a value, or `undefined` if not possible. */
+  resolve(): any | undefined;
+  /**
+   * Evaluate the expression on CPU, returning the result.
+   *
+   * Typically you would compile the AluExp as a representation to a lower-level
+   * language. This is just to define the semantics and help debug.
+   *
+   * Note that the representation of Bool is as a number (0 or 1) here.
+   */
+  evaluate(context: Record<string, any>, globals?: (gid: number, bufidx: number) => any): number;
+  /** Get this expression in debug format as a string. */
+  toString(): string;
+  /** Generic fold() operation with a reducer over the expression tree. */
+  fold<T = void>(reducer: (exp: AluExp, mappedSrc: T[]) => T): T;
+  /** Rewrite the expression recursively using a visitor. */
+  rewrite(visitor: (exp: AluExp) => AluExp | undefined | null): AluExp;
+  /** Collect all nodes that satisfy a predicate. */
+  collect(predicate: (exp: AluExp) => boolean): AluExp[];
 }
 /** Symbolic form for each mathematical operation. */
 declare enum AluOp {
-    Add = "Add",
-    Sub = "Sub",
-    Mul = "Mul",
-    Idiv = "Idiv",
-    Mod = "Mod",
-    Min = "Min",
-    Max = "Max",
-    Sin = "Sin",
-    Cos = "Cos",
-    Exp = "Exp",
-    Log = "Log",
-    Reciprocal = "Reciprocal",
-    Cast = "Cast",
-    Cmplt = "Cmplt",
-    Cmpne = "Cmpne",
-    Where = "Where",// Ternary operator: `cond ? a : b`
-    Const = "Const",// arg = value
-    Special = "Special",// arg = [variable, n]
-    Variable = "Variable",// arg = variable
-    GlobalIndex = "GlobalIndex",// arg = gid; src = [bufidx]
-    GlobalView = "GlobalView"
+  Add = "Add",
+  Sub = "Sub",
+  Mul = "Mul",
+  Idiv = "Idiv",
+  Mod = "Mod",
+  Min = "Min",
+  Max = "Max",
+  Sin = "Sin",
+  Cos = "Cos",
+  Exp = "Exp",
+  Log = "Log",
+  Reciprocal = "Reciprocal",
+  Cast = "Cast",
+  Bitcast = "Bitcast",
+  Cmplt = "Cmplt",
+  Cmpne = "Cmpne",
+  Where = "Where",
+  // Ternary operator: `cond ? a : b`
+  Threefry2x32 = "Threefry2x32",
+  // PRNG operation, arg = 'xor' | 0 | 1
+  Const = "Const",
+  // arg = value
+  Special = "Special",
+  // arg = [variable, n]
+  Variable = "Variable",
+  // arg = variable
+  GlobalIndex = "GlobalIndex",
+  // arg = gid; src = [bufidx]
+  GlobalView = "GlobalView",
 }
 /**
  * Description of a kernel to be compiled.
@@ -231,24 +253,26 @@ declare enum AluOp {
  * indexing into a buffer.
  */
 declare class Kernel implements FpHashable {
-    /** Number of global arguments / arrays. */
-    readonly nargs: number;
-    /** Size of the result array in element count. */
-    readonly size: number;
-    /** Expression to be evaluated. */
-    readonly exp: AluExp;
-    /** Optional reduction to be performed. */
-    readonly reduction?: Reduction | undefined;
-    constructor(
-    /** Number of global arguments / arrays. */
-    nargs: number, 
-    /** Size of the result array in element count. */
-    size: number, 
-    /** Expression to be evaluated. */
-    exp: AluExp, 
-    /** Optional reduction to be performed. */
-    reduction?: Reduction | undefined);
-    hash(state: FpHash): void;
+  /** Number of global arguments / arrays. */
+  readonly nargs: number;
+  /** Size of the result array in element count. */
+  readonly size: number;
+  /** Expression to be evaluated. */
+  readonly exp: AluExp;
+  /** Optional reduction to be performed. */
+  readonly reduction?: Reduction | undefined;
+  constructor(/** Number of global arguments / arrays. */
+  nargs: number, /** Size of the result array in element count. */
+  size: number, /** Expression to be evaluated. */
+  exp: AluExp, /** Optional reduction to be performed. */
+  reduction?: Reduction | undefined);
+  hash(state: FpHash): void;
+  pprint(): PPrint;
+  toString(): string;
+  /** The dtype of the values output by this kernel. */
+  get dtype(): DType;
+  /** The number of bytes in the output array when evaluating this kernel. */
+  get bytes(): number;
 }
 /**
  * Description of a reduction.
@@ -266,41 +290,29 @@ declare class Kernel implements FpHashable {
  * at this level since they depend on GPU, versus CPU or Wasm.
  */
 declare class Reduction implements FpHashable {
-    /** Data type of the values being reduced over. */
-    readonly dtype: DType;
-    /** Operation to perform. Only ops in `AluGroup.Reduce` are supported. */
-    readonly op: AluOp;
-    /** Size of the reduction axis. */
-    readonly size: number;
-    /** Follow-up expression defined with the "acc" variable, defaults to identity. */
-    readonly fusion: AluExp;
-    constructor(
-    /** Data type of the values being reduced over. */
-    dtype: DType, 
-    /** Operation to perform. Only ops in `AluGroup.Reduce` are supported. */
-    op: AluOp, 
-    /** Size of the reduction axis. */
-    size: number, 
-    /** Follow-up expression defined with the "acc" variable, defaults to identity. */
-    fusion?: AluExp);
-    hash(state: FpHash): void;
-    /** Get the identity for this reduction operation. */
-    get identity(): any;
-    /** Evaluate this operation on CPU. */
-    evaluate(...values: any): any;
+  /** Data type of the values being reduced over. */
+  readonly dtype: DType;
+  /** Operation to perform. Only ops in `AluGroup.Reduce` are supported. */
+  readonly op: AluOp;
+  /** Size of the reduction axis. */
+  readonly size: number;
+  /** Follow-up expression defined with the "acc" variable, defaults to identity. */
+  readonly fusion: AluExp;
+  constructor(/** Data type of the values being reduced over. */
+  dtype: DType, /** Operation to perform. Only ops in `AluGroup.Reduce` are supported. */
+  op: AluOp, /** Size of the reduction axis. */
+  size: number, /** Follow-up expression defined with the "acc" variable, defaults to identity. */
+  fusion?: AluExp);
+  hash(state: FpHash): void;
+  toString(): string;
+  /** Get the identity for this reduction operation. */
+  get identity(): any;
+  /** Evaluate this operation on CPU. */
+  evaluate(...values: any): any;
 }
-
-/**
- * @file Shared interfaces and code for the low-level backend API.
- *
- * Think of each backend as a _connector_ to a specific hardware or software
- * implementation of the array API.
- *
- * Backends do not share any of the built-in operational semantics of the
- * library. This is a private API. You must allocate and free buffers manually,
- * and dispatch happens on the level of each shader. Buffers are untyped.
- */
-
+/** Expression for accessing `indices` in input array with the given shape. */
+//#endregion
+//#region src/backend.d.ts
 type Device = "cpu" | "webgpu";
 declare const devices: Device[];
 /** Set the default device backend (must be initialized). */
@@ -313,75 +325,74 @@ declare function setDevice(device: Device): void;
  * available backends.
  */
 declare function init(...devicesToInit: Device[]): Promise<Device[]>;
+/** Retrieve a backend that has been initialized. */
+
 /** Unique identifier for an allocated, on-device buffer. */
 type Slot = number;
 /** A device backend. */
 interface Backend {
-    /** The name of the backend as a string. */
-    readonly type: Device;
-    /** Maximum number of arguments per dispatched kernel. */
-    readonly maxArgs: number;
-    /** Allocate a new slot with reference count 1. */
-    malloc(size: number, initialData?: ArrayBuffer): Slot;
-    /** Increment the reference count of the slot. */
-    incRef(slot: Slot): void;
-    /**
-     * Decrement the reference count of the slot. If the reference count reaches
-     * zero, it is freed. This should throw if the slot was already freed.
-     */
-    decRef(slot: Slot): void;
-    /** Read a range of bytes from a buffer. */
-    read(slot: Slot, start?: number, count?: number): Promise<ArrayBuffer>;
-    /** Read a range of bytes from a buffer, blocking variant. */
-    readSync(slot: Slot, start?: number, count?: number): ArrayBuffer;
-    /** Prepare an expression to be executed later. */
-    prepare(kernel: Kernel): Promise<Executable>;
-    /** Prepare an expression to be executed later, blocking variant. */
-    prepareSync(kernel: Kernel): Executable;
-    /**
-     * Run a backend operation that was previously prepared.
-     *
-     * The operation may not run immediately, but operations are guaranteed to run
-     * in the dispatch order. Also, `read()` will wait for all pending operations
-     * on that slot to finish.
-     */
-    dispatch(exe: Executable, inputs: Slot[], outputs: Slot[]): void;
+  /** The name of the backend as a string. */
+  readonly type: Device;
+  /** Maximum number of arguments per dispatched kernel. */
+  readonly maxArgs: number;
+  /** Allocate a new slot with reference count 1. */
+  malloc(size: number, initialData?: ArrayBuffer): Slot;
+  /** Increment the reference count of the slot. */
+  incRef(slot: Slot): void;
+  /**
+   * Decrement the reference count of the slot. If the reference count reaches
+   * zero, it is freed. This should throw if the slot was already freed.
+   */
+  decRef(slot: Slot): void;
+  /** Read a range of bytes from a buffer. */
+  read(slot: Slot, start?: number, count?: number): Promise<ArrayBuffer>;
+  /** Read a range of bytes from a buffer, blocking variant. */
+  readSync(slot: Slot, start?: number, count?: number): ArrayBuffer;
+  /** Prepare an expression to be executed later. */
+  prepare(kernel: Kernel): Promise<Executable>;
+  /** Prepare an expression to be executed later, blocking variant. */
+  prepareSync(kernel: Kernel): Executable;
+  /**
+   * Run a backend operation that was previously prepared.
+   *
+   * The operation may not run immediately, but operations are guaranteed to run
+   * in the dispatch order. Also, `read()` will wait for all pending operations
+   * on that slot to finish.
+   */
+  dispatch(exe: Executable, inputs: Slot[], outputs: Slot[]): void;
 }
 declare class Executable<T = any> {
-    readonly kernel: Kernel;
-    /** Extra data specific to the backend running this kernel. */
-    readonly data: T;
-    constructor(kernel: Kernel, 
-    /** Extra data specific to the backend running this kernel. */
-    data: T);
+  readonly kernel: Kernel;
+  /** Extra data specific to the backend running this kernel. */
+  readonly data: T;
+  constructor(kernel: Kernel, /** Extra data specific to the backend running this kernel. */
+  data: T);
+}
+declare namespace tree_d_exports {
+  export { JsTree, JsTreeDef, MapJsTree, NodeType, flatten, leaves, map, ref, structure, unflatten };
 }
-
-/** @file Utilities for working with tree-like container data structures ("pytrees"). */
 declare enum NodeType {
-    Array = "Array",
-    Object = "Object",
-    Leaf = "Leaf"
+  Array = "Array",
+  Object = "Object",
+  Leaf = "Leaf",
 }
 type JsTree<T> = T | JsTree<T>[] | {
-    [key: string]: JsTree<T>;
-};
-type MapJsTree<T, A, B> = T extends A ? B : T extends globalThis.Array<infer U> ? number extends T["length"] ? MapJsTree<U, A, B>[] : {
-    [K in keyof T]: MapJsTree<T[K], A, B>;
-} : {
-    [K in keyof T]: MapJsTree<T[K], A, B>;
+  [key: string]: JsTree<T>;
 };
+type MapJsTree<T, A, B> = T extends A ? B : T extends globalThis.Array<infer U> ? number extends T["length"] ? MapJsTree<U, A, B>[] : { [K in keyof T]: MapJsTree<T[K], A, B> } : { [K in keyof T]: MapJsTree<T[K], A, B> };
 /** Analog to the JAX "pytree" object, but for JavaScript. */
 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[]);
-    /** Returns a string representation of this tree definition. */
-    toString(root?: boolean): string;
-    /** Compare this tree definition with another. */
-    equals(other: JsTreeDef): boolean;
+  readonly nodeType: NodeType;
+  readonly nodeMetadata: any;
+  readonly childTreedefs: JsTreeDef[];
+  static leaf: JsTreeDef;
+  constructor(nodeType: NodeType, nodeMetadata: any,
+  // Must be comparable with deepEqual.
+  childTreedefs: JsTreeDef[]);
+  /** 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];
@@ -391,151 +402,290 @@ declare function leaves<T>(tree: JsTree<T>): T[];
 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>;
-
-type tree_JsTree<T> = JsTree<T>;
-type tree_JsTreeDef = JsTreeDef;
-declare const tree_JsTreeDef: typeof JsTreeDef;
-type tree_MapJsTree<T, A, B> = MapJsTree<T, A, B>;
-type tree_NodeType = NodeType;
-declare const tree_NodeType: typeof NodeType;
-declare const tree_flatten: typeof flatten;
-declare const tree_leaves: typeof leaves;
-declare const tree_structure: typeof structure;
-declare const tree_unflatten: typeof unflatten;
-declare namespace tree {
-  export { type tree_JsTree as JsTree, tree_JsTreeDef as JsTreeDef, type tree_MapJsTree as MapJsTree, tree_NodeType as NodeType, tree_flatten as flatten, tree_leaves as leaves, tree_structure as structure, tree_unflatten as unflatten };
-}
-
-/** @file Core library internals and interpreter stack, based on Autodidax. */
-
+/** 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;
+//#endregion
+//#region src/frontend/core.d.ts
+/**
+ * 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).
+ *
+ * All n-ary operations support broadcasting, with NumPy semantics.
+ */
 declare enum Primitive {
-    Add = "add",
-    Mul = "mul",
-    Idiv = "idiv",
-    Neg = "neg",
-    Reciprocal = "reciprocal",
-    Sin = "sin",
-    Cos = "cos",
-    Exp = "exp",
-    Log = "log",
-    Min = "min",
-    Max = "max",
-    ReduceSum = "reduce_sum",
-    Compare = "compare",
-    Where = "where",
-    Transpose = "transpose",
-    Broadcast = "broadcast",
-    Reshape = "reshape",
-    Flip = "flip",
-    JitCall = "jit_call"
+  Add = "add",
+  Mul = "mul",
+  Idiv = "idiv",
+  Neg = "neg",
+  Reciprocal = "reciprocal",
+  StopGradient = "stop_gradient",
+  Cast = "cast",
+  Bitcast = "bitcast",
+  RandomBits = "random_bits",
+  Sin = "sin",
+  Cos = "cos",
+  Exp = "exp",
+  Log = "log",
+  Min = "min",
+  Max = "max",
+  Reduce = "reduce",
+  Dot = "dot",
+  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, unknown>> {
+  [Primitive.Cast]: {
+    dtype: DType;
+  };
+  [Primitive.Bitcast]: {
+    dtype: DType;
+  };
+  [Primitive.Reduce]: {
+    op: AluOp;
+    axis: 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: [number, number][];
+  };
+  [Primitive.Pad]: {
+    width: [number, number][];
+  };
+  [Primitive.Gather]: {
+    axis: number[];
+    outDim: number;
+  };
+  [Primitive.JitCall]: {
+    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 {
+  Greater = "greater",
+  Less = "less",
+  Equal = "equal",
+  NotEqual = "not_equal",
+  GreaterEqual = "greater_equal",
+  LessEqual = "less_equal",
 }
+/** @inline */
+type ReduceOpts = {
+  keepDims?: boolean;
+};
 type MainTrace = {
-    level: number;
-    traceType: new (main: MainTrace) => Trace;
-    globalData: any | null;
+  level: number;
+  traceType: new (main: MainTrace) => Trace;
+  globalData: any | null;
 };
+/**
+ * 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(primitive: Primitive, tracers: Tracer[], params: Record<string, any>): Tracer[];
+  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[];
 }
 interface AbstractValue {
-    shape: number[];
-    dtype: DType;
+  shape: number[];
+  dtype: DType;
 }
 declare abstract class Tracer {
-    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;
-    get shape(): number[];
-    get dtype(): DType;
-    get ndim(): number;
-    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(axis?: number | number[]): this;
-    transpose(perm?: number[]): this;
-    reshape(shape: number | number[]): 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;
+  /** @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;
+  get shape(): number[];
+  get dtype(): DType;
+  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?: number | number[], opts?: ReduceOpts): this;
+  /** Product of the array elements over a given axis. */
+  prod(axis?: number | number[], opts?: ReduceOpts): this;
+  /** Compute the average of the array elements along the specified axis. */
+  mean(axis?: number | number[], 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] | [number, number] | null | Tracer)[]): this;
 }
 declare class ShapedArray implements AbstractValue {
-    readonly shape: number[];
-    readonly dtype: DType;
-    constructor(shape: number[], dtype: DType);
-    static fromAval(aval: AbstractValue): ShapedArray;
-    get ndim(): number;
-    strShort(): string;
-    equals(other: ShapedArray): boolean;
+  readonly shape: number[];
+  readonly dtype: DType;
+  constructor(shape: number[], dtype: DType);
+  static fromAval(aval: AbstractValue): ShapedArray;
+  get ndim(): number;
+  strShort(): 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.
  *
@@ -543,22 +693,23 @@ type ArrayLike = Array | number | boolean;
  * 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;
+  #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;
+  dtype?: DType;
+  device?: Device;
 };
 /**
  * A multidimensional numeric array with data stored on CPU or GPU.
@@ -571,64 +722,104 @@ type DTypeAndDevice = {
  * "Array" type by name.
  */
 declare class Array extends Tracer {
-    #private;
-    id: number;
-    constructor(source: AluExp | Slot, st: ShapeTracker, dtype: DType, backend: Backend, pending?: Iterable<PendingExecute> | null);
-    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<Float32Array | Int32Array>;
-    /** Wait for this array to be placed on the backend, if needed. */
-    wait(): Promise<void>;
-    /**
-     * 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(): Float32Array | Int32Array;
-    /** Convert this array into a JavaScript object (blocking). */
-    js(): RecursiveArray<number> | RecursiveArray<boolean>;
-    /** Convert this array into a JavaScript object, asynchronously. */
-    jsAsync(): Promise<RecursiveArray<number> | RecursiveArray<boolean>>;
-    /** @private Internal plumbing method for Array / Tracer ops. */
-    static _implRules(): Record<Primitive, ImplRule>;
-    _realizeSource(): number;
+  #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(source: AluExp | Slot, st: ShapeTracker, dtype: DType, backend: Backend, pending?: Iterable<PendingExecute> | null);
+  /** @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<Float32Array | Int32Array | Uint32Array>;
+  /** Wait for this array to be placed on the backend, if needed. */
+  wait(): Promise<void>;
+  /**
+   * 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(): Float32Array | Int32Array | Uint32Array;
+  /**
+   * 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>;
+  /** @private Internal plumbing method for Array / Tracer ops. */
+  static _implRules(): typeof implRules;
+  _realizeSource(): number;
 }
 /** Construct an array from a single scalar constant. */
-declare function scalar(value: number | boolean, { dtype, device }?: DTypeAndDevice): Array;
+declare function scalar(value: number | boolean, {
+  dtype,
+  device
+}?: DTypeAndDevice): Array;
 /** Constructor for creating a new array from data. */
-declare function array(values: Array | Float32Array | Int32Array | RecursiveArray<number> | RecursiveArray<boolean>, { shape, dtype, device }?: {
-    shape?: number[];
+declare function array(values: Array | Float32Array | Int32Array | RecursiveArray<number> | RecursiveArray<boolean>, {
+  shape,
+  dtype,
+  device
+}?: {
+  shape?: number[];
 } & DTypeAndDevice): Array;
-type ImplRule = (tracers: Array[], params: any) => 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;
+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;
+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;
+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;
+declare function eye(numRows: number, numCols?: number, {
+  dtype,
+  device
+}?: DTypeAndDevice): Array;
 /** Return the identity array, with ones on the main diagonal. */
-declare function identity$1(n: number, { dtype, device }?: DTypeAndDevice): Array;
+declare function identity$1(n: number, {
+  dtype,
+  device
+}?: DTypeAndDevice): Array;
 /**
  * Return evenly spaced values within a given interval.
  *
@@ -643,7 +834,10 @@ declare function identity$1(n: number, { dtype, device }?: DTypeAndDevice): Arra
  * 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;
+declare function arange(start: number, stop?: number, step?: number, {
+  dtype,
+  device
+}?: DTypeAndDevice): Array;
 /**
  * Return evenly spaced numbers over a specified interval.
  *
@@ -653,10 +847,17 @@ declare function arange(start: number, stop?: number, step?: number, { dtype, de
  *
  * 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 function linspace(start: number, stop: number, num?: number, endpoint?: boolean, {
+  dtype,
+  device
+}?: DTypeAndDevice): Array;
+/** Translate a `CompareOp` into an `AluExp` on two sub-expressions. */
+declare namespace numpy_d_exports {
+  export { Array, ArrayLike, DType, abs, absolute, add, allclose, arange, argmax, argmin, array, astype, bool, clip, columnStack, complex64, concatenate, cos, diag, diagonal, divide, dot, dstack, e, equal, eulerGamma, exp, exp2, eye, flip, fliplr, flipud, float32, full, greater, greaterEqual, hstack, identity$1 as identity, inf, int32, less, lessEqual, linspace, log, log10, log2, matmul, max, maximum, mean, meshgrid, min, minimum, moveaxis, multiply, nan, ndim, negative, notEqual, ones, pad, permuteDims, pi, prod, ravel, reciprocal, reshape, scalar, shape$1 as shape, sin, size, square, stack, sum, tan, transpose, trueDivide, trunc, uint32, vdot, vecdot, vstack, where, zeros };
+}
 declare const float32 = DType.Float32;
 declare const int32 = DType.Int32;
+declare const uint32 = DType.Uint32;
 declare const bool = DType.Bool;
 declare const complex64 = DType.Complex64;
 /** Euler's constant, `e = 2.7182818284590...` */
@@ -712,16 +913,92 @@ declare const transpose: (x: ArrayLike, perm?: number[]) => Array;
  * length of the array and remaining dimensions.
  */
 declare const reshape: (x: ArrayLike, shape: number[]) => Array;
-declare const sum: (x: ArrayLike, axis?: number | number[]) => Array;
+/** Move axes of an array to new positions. Other axes retain original order. */
 declare const moveaxis: (x: ArrayLike, src: number, dst: number) => Array;
-/** Return the number of dimensions of an array. */
+/**
+ * Add padding (zeros) to an array.
+ *
+ * The `width` argument is either an integer or pair of integers, in which case
+ * all axes are padded with the same width. Or if it is an array of pairs, each
+ * pair specifies the padding for its corresponding axis.
+ */
+declare const pad: (x: ArrayLike, width: number | [number, number] | [number, number][]) => Array;
+/** Return the number of dimensions of an array. Does not consume array reference. */
 declare const ndim: (x: ArrayLike) => number;
-/** Return the shape of an array. */
-declare const shape: (x: ArrayLike) => number[];
-/** Return the number of elements in an array, optionally along an axis. */
+/** Return the shape of an array. Does not consume array reference. */
+declare const shape$1: (x: ArrayLike) => number[];
+/**
+ * Return the number of elements in an array, optionally along an axis.
+ * Does not consume array reference.
+ */
 declare function size(a: ArrayLike, axis?: number): number;
+/** Convert an array to a specified dtype. */
+declare function astype(a: ArrayLike, dtype: DType): Array;
+/** Sum of the elements of the array over a given axis, or axes. */
+declare function sum(a: ArrayLike, axis?: number | number[], opts?: ReduceOpts): Array;
+/** Product of the array elements over a given axis. */
+declare function prod(a: ArrayLike, axis?: number | number[], opts?: ReduceOpts): Array;
+/** Return the minimum of array elements along a given axis. */
+declare function min(a: ArrayLike, axis?: number | number[], opts?: ReduceOpts): Array;
+/** Return the maximum of array elements along a given axis. */
+declare function max(a: ArrayLike, axis?: number | number[], opts?: ReduceOpts): Array;
+/** Compute the average of the array elements along the specified axis. */
+declare function mean(a: ArrayLike, axis?: number | number[], opts?: ReduceOpts): Array;
+/**
+ * Returns the indices of the minimum values along an axis.
+ *
+ * By default, index is into the flatted array, otherwise it is along the
+ * specified axis.
+ */
+declare function argmin(a: ArrayLike, axis?: number, opts?: ReduceOpts): Array;
+/**
+ * Returns the indices of the maximum values along an axis.
+ *
+ * By default, index is into the flatted array, otherwise it is along the
+ * specified axis.
+ */
+declare function argmax(a: ArrayLike, axis?: number, opts?: ReduceOpts): Array;
 /** Reverse the elements in an array along the given axes. */
 declare function flip(x: ArrayLike, axis?: number | number[]): Array;
+/**
+ * Join a sequence of arrays along an existing axis.
+ *
+ * The arrays must have the same shape, except in the dimension corresponding to
+ * `axis` (the first, by default).
+ *
+ * No scalars can be passed to this function, as the axis is then ambiguous.
+ */
+declare function concatenate(xs: Array[], axis?: number): Array;
+/**
+ * Join a sequence of arrays along a new axis.
+ *
+ * The `axis` parameter specifies the index of the new axis in the dimensions of
+ * the result. For example, if `axis=0` it will be the first dimension and if
+ * `axis=-1` it will be the last dimension.
+ *
+ * All shapes must have the same shape.
+ */
+declare function stack(xs: ArrayLike[], axis?: number): Array;
+/**
+ * Horizontally stack arrays. Inputs are promoted to rank at least 1, then
+ * concatenated along axis 1 (if rank-2 or higher) or 0 (if rank-1).
+ */
+declare function hstack(xs: ArrayLike[]): Array;
+/**
+ * Vertically stack arrays. Inputs are promoted to rank at least 2, then
+ * concatenated along axis 0.
+ */
+declare function vstack(xs: ArrayLike[]): Array;
+/**
+ * Stack arrays depth-wise. Inputs are promoted to rank at least 3, then
+ * concatenated along axis 2.
+ */
+declare function dstack(xs: ArrayLike[]): Array;
+/**
+ * Stack arrays column-wise. Inputs are promoted to rank at least 2, then
+ * concatenated along axis 1.
+ */
+declare function columnStack(xs: ArrayLike[]): Array;
 /** Flip an array vertically (axis=0). */
 declare function flipud(x: ArrayLike): Array;
 /** Flip an array horizontally (axis=1). */
@@ -738,8 +1015,6 @@ declare function ravel(a: ArrayLike): Array;
  * This returns a view over the existing array.
  */
 declare function diagonal(a: ArrayLike, offset?: number, axis1?: number, axis2?: number): Array;
-/** Transposes a matrix or stack of matrices `x` (swap last two axes). */
-declare function matrixTranspose(x: ArrayLike): Array;
 /**
  * Extract a diagonal or construct a diagonal array.
  *
@@ -749,15 +1024,15 @@ declare function matrixTranspose(x: ArrayLike): Array;
 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;
+  rtol?: number;
+  atol?: number;
 }): boolean;
 /** Matrix product of two arrays. */
-declare const matmul: (x: ArrayLike, y: ArrayLike) => Array;
+declare function matmul(x: ArrayLike, y: ArrayLike): Array;
 /** Dot product of two arrays. */
-declare const dot: (x: ArrayLike, y: ArrayLike) => Array;
+declare function dot(x: ArrayLike, y: ArrayLike): Array;
 /** Vector dot product of two arrays. */
-declare const vecdot: (x: ArrayLike, y: ArrayLike) => Array;
+declare function vecdot(x: ArrayLike, y: ArrayLike): Array;
 /**
  * Return the dot product of two vectors.
  *
@@ -770,8 +1045,10 @@ declare function vdot(x: ArrayLike, y: ArrayLike): 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";
+declare function meshgrid(xs: Array[], {
+  indexing
+}?: {
+  indexing?: "xy" | "ij";
 }): Array[];
 /**
  * Clip (limit) the values in an array.
@@ -807,161 +1084,71 @@ declare function exp2(p: ArrayLike): Array;
 declare function log2(x: ArrayLike): Array;
 /** Return the base-10 logarithm of x, element-wise. */
 declare function log10(x: ArrayLike): Array;
-
-type numpy_Array = Array;
-declare const numpy_Array: typeof Array;
-type numpy_ArrayLike = ArrayLike;
-type numpy_DType = DType;
-declare const numpy_DType: typeof DType;
-declare const numpy_abs: typeof abs;
-declare const numpy_absolute: typeof absolute;
-declare const numpy_add: typeof add;
-declare const numpy_allclose: typeof allclose;
-declare const numpy_arange: typeof arange;
-declare const numpy_array: typeof array;
-declare const numpy_bool: typeof bool;
-declare const numpy_clip: typeof clip;
-declare const numpy_complex64: typeof complex64;
-declare const numpy_cos: typeof cos;
-declare const numpy_diag: typeof diag;
-declare const numpy_diagonal: typeof diagonal;
-declare const numpy_divide: typeof divide;
-declare const numpy_dot: typeof dot;
-declare const numpy_e: typeof e;
-declare const numpy_equal: typeof equal;
-declare const numpy_eulerGamma: typeof eulerGamma;
-declare const numpy_exp: typeof exp;
-declare const numpy_exp2: typeof exp2;
-declare const numpy_eye: typeof eye;
-declare const numpy_flip: typeof flip;
-declare const numpy_fliplr: typeof fliplr;
-declare const numpy_flipud: typeof flipud;
-declare const numpy_float32: typeof float32;
-declare const numpy_full: typeof full;
-declare const numpy_greater: typeof greater;
-declare const numpy_greaterEqual: typeof greaterEqual;
-declare const numpy_inf: typeof inf;
-declare const numpy_int32: typeof int32;
-declare const numpy_less: typeof less;
-declare const numpy_lessEqual: typeof lessEqual;
-declare const numpy_linspace: typeof linspace;
-declare const numpy_log: typeof log;
-declare const numpy_log10: typeof log10;
-declare const numpy_log2: typeof log2;
-declare const numpy_matmul: typeof matmul;
-declare const numpy_matrixTranspose: typeof matrixTranspose;
-declare const numpy_maximum: typeof maximum;
-declare const numpy_meshgrid: typeof meshgrid;
-declare const numpy_minimum: typeof minimum;
-declare const numpy_moveaxis: typeof moveaxis;
-declare const numpy_multiply: typeof multiply;
-declare const numpy_nan: typeof nan;
-declare const numpy_ndim: typeof ndim;
-declare const numpy_negative: typeof negative;
-declare const numpy_notEqual: typeof notEqual;
-declare const numpy_ones: typeof ones;
-declare const numpy_permuteDims: typeof permuteDims;
-declare const numpy_pi: typeof pi;
-declare const numpy_ravel: typeof ravel;
-declare const numpy_reciprocal: typeof reciprocal;
-declare const numpy_reshape: typeof reshape;
-declare const numpy_scalar: typeof scalar;
-declare const numpy_shape: typeof shape;
-declare const numpy_sin: typeof sin;
-declare const numpy_size: typeof size;
-declare const numpy_square: typeof square;
-declare const numpy_sum: typeof sum;
-declare const numpy_tan: typeof tan;
-declare const numpy_transpose: typeof transpose;
-declare const numpy_trueDivide: typeof trueDivide;
-declare const numpy_trunc: typeof trunc;
-declare const numpy_vdot: typeof vdot;
-declare const numpy_vecdot: typeof vecdot;
-declare const numpy_where: typeof where;
-declare const numpy_zeros: typeof zeros;
-declare namespace numpy {
-  export { numpy_Array as Array, type numpy_ArrayLike as ArrayLike, numpy_DType as DType, numpy_abs as abs, numpy_absolute as absolute, numpy_add as add, numpy_allclose as allclose, numpy_arange as arange, numpy_array as array, numpy_bool as bool, numpy_clip as clip, numpy_complex64 as complex64, numpy_cos as cos, numpy_diag as diag, numpy_diagonal as diagonal, numpy_divide as divide, numpy_dot as dot, numpy_e as e, numpy_equal as equal, numpy_eulerGamma as eulerGamma, numpy_exp as exp, numpy_exp2 as exp2, numpy_eye as eye, numpy_flip as flip, numpy_fliplr as fliplr, numpy_flipud as flipud, numpy_float32 as float32, numpy_full as full, numpy_greater as greater, numpy_greaterEqual as greaterEqual, identity$1 as identity, numpy_inf as inf, numpy_int32 as int32, numpy_less as less, numpy_lessEqual as lessEqual, numpy_linspace as linspace, numpy_log as log, numpy_log10 as log10, numpy_log2 as log2, numpy_matmul as matmul, numpy_matrixTranspose as matrixTranspose, numpy_maximum as maximum, numpy_meshgrid as meshgrid, numpy_minimum as minimum, numpy_moveaxis as moveaxis, numpy_multiply as multiply, numpy_nan as nan, numpy_ndim as ndim, numpy_negative as negative, numpy_notEqual as notEqual, numpy_ones as ones, numpy_permuteDims as permuteDims, numpy_pi as pi, numpy_ravel as ravel, numpy_reciprocal as reciprocal, numpy_reshape as reshape, numpy_scalar as scalar, numpy_shape as shape, numpy_sin as sin, numpy_size as size, numpy_square as square, numpy_sum as sum, numpy_tan as tan, numpy_transpose as transpose, numpy_trueDivide as trueDivide, numpy_trunc as trunc, numpy_vdot as vdot, numpy_vecdot as vecdot, numpy_where as where, numpy_zeros as zeros };
-}
-
-/** General class for pretty-printing expressions with indentation. */
-declare class PPrint {
-    readonly indents: number[];
-    readonly lines: string[];
-    constructor(indents: number[], lines: string[]);
-    /** Add a fixed amount of indentation to each line. */
-    indent(spaces: number): PPrint;
-    /** Concatenate two or more pretty-printed expressions. */
-    concat(...items: PPrint[]): PPrint;
-    /** Stack one block to the right of another one, sharing 1 common line. */
-    stack(other: PPrint): PPrint;
-    /** Combine this block of lines into a formatted string. */
-    toString(): string;
-    static pp(s: Stringable): PPrint;
-}
-interface Stringable {
-    toString(): string;
-}
-
+//#endregion
+//#region src/frontend/jaxpr.d.ts
 /** Variable in a Jaxpr expression. */
 declare class Var {
-    #private;
-    readonly id: number;
-    readonly aval: ShapedArray;
-    constructor(aval: ShapedArray);
-    toString(): string;
+  #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 dtype: DType;
-    readonly value: number;
-    readonly aval: ShapedArray;
-    constructor(dtype: DType, value: number);
+  readonly dtype: DType;
+  readonly value: number;
+  readonly aval: ShapedArray;
+  constructor(dtype: DType, value: number);
 }
 type Atom = Var | Lit;
 declare class VarPrinter {
-    #private;
-    names: Map<Var, string>;
-    name(v: Var): string;
-    nameType(v: Var): string;
+  #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;
+  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.
-     */
-    simplify(): Jaxpr;
-    /** Flattens nested JitCall in a Jaxpr. Useful for handling jit-of-jit. */
-    flatten(): Jaxpr;
+  #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;
+}
+declare namespace nn_d_exports {
+  export { identity, logSigmoid, logSoftmax, logsumexp, oneHot, relu, relu6, sigmoid, silu, softSign, softmax, softplus, swish };
 }
-
 /**
  * Rectified Linear Unit (ReLU) activation function:
  * `relu(x) = max(x, 0)`.
@@ -1018,20 +1205,70 @@ declare const swish: typeof silu;
 declare function logSigmoid(x: ArrayLike): Array;
 /** Identity activation function. Returns the argument unmodified. */
 declare const identity: (x: ArrayLike) => Array;
-
-declare const nn_identity: typeof identity;
-declare const nn_logSigmoid: typeof logSigmoid;
-declare const nn_relu: typeof relu;
-declare const nn_relu6: typeof relu6;
-declare const nn_sigmoid: typeof sigmoid;
-declare const nn_silu: typeof silu;
-declare const nn_softSign: typeof softSign;
-declare const nn_softplus: typeof softplus;
-declare const nn_swish: typeof swish;
-declare namespace nn {
-  export { nn_identity as identity, nn_logSigmoid as logSigmoid, nn_relu as relu, nn_relu6 as relu6, nn_sigmoid as sigmoid, nn_silu as silu, nn_softSign as softSign, nn_softplus as softplus, nn_swish as swish };
+/**
+ * Softmax function. Computes the function which rescales elements to the range
+ * [0, 1] such that the elements along `axis` sum to 1.
+ *
+ * If `axis` is not specified, it defaults to the last axis.
+ *
+ * Reference: https://en.wikipedia.org/wiki/Softmax_function
+ */
+declare function softmax(x: ArrayLike, axis?: number | number[]): Array;
+/**
+ * Log-Softmax function.
+ *
+ * Computes the logarithm of the `softmax` function, which rescales elements to
+ * the range [-infinity, 0).
+ *
+ * If `axis` is not specified, it defaults to the last axis.
+ */
+declare function logSoftmax(x: ArrayLike, axis?: number | number[]): Array;
+/**
+ * Log-sum-exp reduction. Also a multivariate version of `softplus`.
+ *
+ * If no axis is specified, the reduction is performed over all elements. This
+ * convention differs from `jax.nn.logSoftmax()`.
+ *
+ * Reference: https://en.wikipedia.org/wiki/LogSumExp
+ */
+declare function logsumexp(x: ArrayLike, axis?: number | number[]): Array;
+/**
+ * One-hot encodes the given indices.
+ *
+ * Each index in the integer input `x` is encoded as a vector of zeros of length
+ * `numClasses`, with a 1 at the index position specified by its value.
+ *
+ * ```js
+ * import { nn, numpy as np } from '@jax-js/jax';
+ *
+ * nn.oneHot(np.array([1, 1, 2], { dtype: np.int32 }), 3);
+ * // Output:
+ * // [[0, 1, 0],
+ * //  [0, 1, 0],
+ * //  [0, 0, 1]]
+ * ```
+ */
+declare function oneHot(x: Array, numClasses: number): Array;
+declare namespace random_d_exports {
+  export { bits, key, split, uniform };
 }
-
+/** Create a pseudo-random number generator (PRNG) key from 32-bit integer seed. */
+declare function key(seed: number): Array;
+/** Splits a PRNG key into `num` new keys by adding a leading axis. */
+declare function split(key: Array, num?: number | number[]): Array;
+/** Sample uniform bits in the form of unsigned integers. */
+declare function bits(key: Array, shape?: number[]): Array;
+/** Sample uniform random values in [minval, maxval) with given shape. */
+declare function uniform(key: Array, shape?: number[], {
+  minval,
+  maxval
+}?: {
+  minval?: number;
+  maxval?: number;
+}): Array;
+//#endregion
+//#region src/index.d.ts
+/** @inline */
 type WithArgsSubtype<F extends (args: any[]) => any, T> = Parameters<F> extends T ? F : never;
 /** Compute the forward-mode Jacobian-vector product for a function. */
 declare const jvp: <F extends (...args: any[]) => JsTree<Array>>(f: WithArgsSubtype<F, JsTree<ArrayLike>>, primals: MapJsTree<Parameters<F>, ArrayLike, ArrayLike>, tangents: MapJsTree<Parameters<F>, ArrayLike, ArrayLike>) => [ReturnType<F>, ReturnType<F>];
@@ -1041,9 +1278,9 @@ declare const vmap: <F extends (...args: any[]) => JsTree<Array>>(f: WithArgsSub
 declare const jacfwd: <F extends (x: Array) => Array>(f: F) => (...args: MapJsTree<Parameters<F>, ArrayLike, ArrayLike>) => ReturnType<F>;
 /** Construct a Jaxpr by dynamically tracing a function with example inputs. */
 declare const makeJaxpr: <F extends (...args: any[]) => JsTree<Array>>(f: WithArgsSubtype<F, JsTree<ArrayLike>>) => (...args: Parameters<F>) => {
-    jaxpr: Jaxpr;
-    consts: Array[];
-    treedef: JsTreeDef;
+  jaxpr: Jaxpr;
+  consts: Array[];
+  treedef: JsTreeDef;
 };
 declare const jit: <F extends (...args: any[]) => JsTree<Array>>(f: WithArgsSubtype<F, JsTree<ArrayLike>>) => (...args: MapJsTree<Parameters<F>, ArrayLike, ArrayLike>) => ReturnType<F>;
 /**
@@ -1058,9 +1295,11 @@ declare const vjp: <F extends (...args: any[]) => JsTree<Array>>(f: WithArgsSubt
  * first argument.
  */
 declare const grad: <F extends (...args: any[]) => JsTree<Array>>(f: WithArgsSubtype<F, JsTree<ArrayLike>>) => (...primals: MapJsTree<Parameters<F>, ArrayLike, ArrayLike>) => MapJsTree<Parameters<F>[0], ArrayLike, Array>;
+/** Create a function that evaluates both `f` and the gradient of `f`. */
+declare const valueAndGrad: <F extends (...args: any[]) => JsTree<Array>>(f: WithArgsSubtype<F, JsTree<ArrayLike>>) => (...primals: MapJsTree<Parameters<F>, ArrayLike, ArrayLike>) => [ReturnType<F>, MapJsTree<Parameters<F>[0], ArrayLike, Array>];
 /** Compute the Jacobian evaluated row-by-row by reverse-mode AD. */
 declare const jacrev: typeof jacfwd;
 /** Compute the Jacobian with reverse-mode AD. Alias for `jacrev()`. */
 declare const jacobian: <F extends (x: Array) => Array>(f: F) => (...args: MapJsTree<Parameters<F>, ArrayLike, ArrayLike>) => ReturnType<F>;
-
-export { type Device, devices, grad, init, jacfwd, jacobian, jacrev, jit, jvp, linearize, makeJaxpr, nn, numpy, setDevice, tree, vjp, vmap };
+//#endregion
+export { type Device, devices, grad, init, jacfwd, jacobian, jacrev, jit, jvp, linearize, makeJaxpr, nn_d_exports as nn, numpy_d_exports as numpy, random_d_exports as random, setDevice, tree_d_exports as tree, valueAndGrad, vjp, vmap };