@jax-js/jax 0.0.1

package/dist/index.d.cts

@@ -0,0 +1,1066 @@
+/**
+ * @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.
+ */
+
+type Pair = [number, number];
+/**
+ * A multidimensional view into memory. An array can be thought of as the
+ * combination of a linear buffer of memory, along with a `View`.
+ *
+ * Formula for getting a data point is basically:
+ *   1. Check if ∀i. 0 <= dim[i] < shape[i], otherwise out of bounds.
+ *   2. If mask exists, and ∃i. dim[i] ∉ mask[i], return 0.
+ *   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;
+}
+/**
+ * Array shape after applying movement operations, as a series of views.
+ *
+ * Each view is applied, then treated as if it were a contiguous array of its
+ * 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;
+}
+
+type RecursiveArray<T> = T | RecursiveArray<T>[];
+interface FpHashable {
+    hash(state: FpHash): void;
+}
+/**
+ * Polynomial hashes modulo p are good at avoiding collisions in expectation.
+ * Probability-wise, it's good enough to be used for something like
+ * deduplicating seen compiler expressions, although it's not adversarial.
+ *
+ * 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;
+}
+
+declare enum DType {
+    Float32 = "float32",
+    Int32 = "int32",
+    Bool = "bool",
+    Complex64 = "complex64"
+}
+/**
+ * Mathematical expression on scalar values.
+ *
+ * This is similiar to and based on tinygrad's UOp class, but it's more specific
+ * to just math on scalars. We're doing this to avoid the complexity of a full
+ * 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[];
+}
+/** 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"
+}
+/**
+ * Description of a kernel to be compiled.
+ *
+ * Each of these can be processed by a backend into some lower-level
+ * representation. It consists of one or more fused operations, optionally
+ * 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;
+}
+/**
+ * Description of a reduction.
+ *
+ * The strategy of jax-js backends is to either handle a standard operation that
+ * is dispatched in a vectorized way over an array, or to reduce over one axis
+ * of some computation. This is a description of the reduction.
+ *
+ * Reduction only supports a few operations, and only over one axis. Users can
+ * always `flatten()` the array before reducing if needed.
+ *
+ * The backend is responsible for implementing the reduction in a way that
+ * minimizes the number of global memory loads, for efficiency. This involves
+ * passing through some optimization strategy. But optimizations are not coded
+ * 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;
+}
+
+/**
+ * @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.
+ */
+
+type Device = "cpu" | "webgpu";
+declare const devices: Device[];
+/** Set the default device backend (must be initialized). */
+declare function setDevice(device: Device): void;
+/**
+ * Initialize `jax-js` library backends.
+ *
+ * By default, this will initialize all available backends. If one or more
+ * backends is provided, only attempt to initialize those. Returns a list of
+ * available backends.
+ */
+declare function init(...devicesToInit: Device[]): Promise<Device[]>;
+/** 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;
+}
+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);
+}
+
+/** @file Utilities for working with tree-like container data structures ("pytrees"). */
+declare enum NodeType {
+    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>;
+};
+/** 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;
+}
+/** 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>;
+
+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. */
+
+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"
+}
+type MainTrace = {
+    level: number;
+    traceType: new (main: MainTrace) => Trace;
+    globalData: any | null;
+};
+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[];
+}
+interface AbstractValue {
+    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;
+}
+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;
+}
+
+type ArrayLike = Array | number | boolean;
+/**
+ * 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;
+}
+type DTypeAndDevice = {
+    dtype?: DType;
+    device?: Device;
+};
+/**
+ * A multidimensional numeric array with data stored on CPU or GPU.
+ *
+ * This is the library's core data type. Equivalent to `jnp.Array` from JAX, or
+ * `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;
+    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;
+}
+/** Construct an array from a single scalar constant. */
+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[];
+} & DTypeAndDevice): Array;
+type ImplRule = (tracers: Array[], params: any) => Array[];
+/** 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 array, 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 const float32 = DType.Float32;
+declare const int32 = DType.Int32;
+declare const bool = DType.Bool;
+declare const complex64 = DType.Complex64;
+/** 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;
+/** Element-wise addition, with broadcasting. */
+declare const add: (x: ArrayLike, y: ArrayLike) => Array;
+/** Element-wise multiplication, with broadcasting. */
+declare const multiply: (x: ArrayLike, y: ArrayLike) => Array;
+/** Numerical negative of every element of an array. */
+declare const negative: (x: ArrayLike) => Array;
+/** Calculate element-wise reciprocal of the input. This is `1/x`. */
+declare const reciprocal: (x: ArrayLike) => Array;
+/** Element-wise sine function (takes radians). */
+declare const sin: (x: ArrayLike) => Array;
+/** Element-wise cosine function (takes radians). */
+declare const cos: (x: ArrayLike) => Array;
+/** Calculate the exponential of all elements in the input array. */
+declare const exp: (x: ArrayLike) => Array;
+/** Calculate the natural logarithm of all elements in the input array. */
+declare const log: (x: ArrayLike) => Array;
+/** Return element-wise minimum of the input arrays. */
+declare const minimum: (x: ArrayLike, y: ArrayLike) => Array;
+/** Return element-wise maximum of the input arrays. */
+declare const maximum: (x: ArrayLike, y: ArrayLike) => Array;
+/** Compare two arrays element-wise. */
+declare const greater: (x: ArrayLike, y: ArrayLike) => Array;
+/** Compare two arrays element-wise. */
+declare const less: (x: ArrayLike, y: ArrayLike) => Array;
+/** Compare two arrays element-wise. */
+declare const equal: (x: ArrayLike, y: ArrayLike) => Array;
+/** Compare two arrays element-wise. */
+declare const notEqual: (x: ArrayLike, y: ArrayLike) => Array;
+/** Compare two arrays element-wise. */
+declare const greaterEqual: (x: ArrayLike, y: ArrayLike) => Array;
+/** Compare two arrays element-wise. */
+declare const lessEqual: (x: ArrayLike, y: ArrayLike) => Array;
+/** Element-wise ternary operator, evaluates to `x` if cond else `y`. */
+declare const where: (cond: ArrayLike, x: ArrayLike, y: ArrayLike) => Array;
+/** Permute the dimensions of an array. Defaults to reversing the axis order. */
+declare const transpose: (x: ArrayLike, perm?: number[]) => Array;
+/**
+ * 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;
+declare const sum: (x: ArrayLike, axis?: number | number[]) => Array;
+declare const moveaxis: (x: ArrayLike, src: number, dst: number) => Array;
+/** Return the number of dimensions of an array. */
+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. */
+declare function size(a: ArrayLike, axis?: number): number;
+/** Reverse the elements in an array along the given axes. */
+declare function flip(x: ArrayLike, axis?: number | number[]): Array;
+/** Flip an array vertically (axis=0). */
+declare function flipud(x: ArrayLike): Array;
+/** Flip an array horizontally (axis=1). */
+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;
+/**
+ * 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.
+ *
+ * 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.
+ *
+ * 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;
+/** 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 const matmul: (x: ArrayLike, y: ArrayLike) => Array;
+/** Dot product of two arrays. */
+declare const dot: (x: ArrayLike, y: ArrayLike) => Array;
+/** Vector dot product of two arrays. */
+declare const vecdot: (x: ArrayLike, y: ArrayLike) => 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[];
+/**
+ * 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;
+/** Alias of `jax.numpy.absolute()`. */
+declare const abs: typeof absolute;
+/** Calculate element-wise square of the input array. */
+declare function square(x: ArrayLike): Array;
+/** Compute a trigonometric tangent of each element of input. */
+declare function tan(x: ArrayLike): Array;
+/** Calculates the floating-point division of x by y element-wise. */
+declare function trueDivide(x: ArrayLike, y: ArrayLike): Array;
+/** 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;
+
+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;
+}
+
+/** 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 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;
+}
+/** 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.
+     */
+    simplify(): Jaxpr;
+    /** Flattens nested JitCall in a Jaxpr. Useful for handling jit-of-jit. */
+    flatten(): Jaxpr;
+}
+
+/**
+ * Rectified Linear Unit (ReLU) activation function:
+ * `relu(x) = max(x, 0)`.
+ */
+declare function relu(x: ArrayLike): Array;
+/**
+ * Rectified Linear Unit 6 (ReLU6) activation function:
+ * `relu6(x) = min(max(x, 0), 6)`.
+ */
+declare function relu6(x: ArrayLike): Array;
+/**
+ * Sigmoid activation function, computed element-wise:
+ * `sigmoid(x) = 1 / (1 + exp(-x))`.
+ *
+ * Reference: https://en.wikipedia.org/wiki/Sigmoid_function
+ */
+declare function sigmoid(x: ArrayLike): Array;
+/**
+ * Softplus activation function:
+ * `softplus(x) = log(1 + exp(x))`.
+ *
+ * Reference: https://en.wikipedia.org/wiki/Softplus
+ */
+declare function softplus(x: ArrayLike): Array;
+/**
+ * Soft-sign activation function, computed element-wise:
+ * `softsign(x) = x / (|x| + 1)`.
+ */
+declare function softSign(x: ArrayLike): Array;
+/**
+ * Sigmoid-weighted Linear Unit (SiLU) activation function, also known as
+ * Swish, computed element-wise:
+ * `silu(x) = x * sigmoid(x) = x / (1 + exp(-x))`.
+ *
+ * `swish()` and `silu()` are both aliases for the same function.
+ *
+ * Reference: https://en.wikipedia.org/wiki/Swish_function
+ */
+declare function silu(x: ArrayLike): Array;
+/**
+ * Sigmoid-weighted Linear Unit (SiLU) activation function, also known as
+ * Swish, computed element-wise:
+ * `silu(x) = x * sigmoid(x) = x / (1 + exp(-x))`.
+ *
+ * `swish()` and `silu()` are both aliases for the same function.
+ *
+ * Reference: https://en.wikipedia.org/wiki/Swish_function
+ */
+declare const swish: typeof silu;
+/**
+ * Log-sigmoid activation function, computed element-wise:
+ * `log_sigmoid(x) = log(sigmoid(x)) = -log(1 + exp(-x))`.
+ */
+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 };
+}
+
+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>];
+/** Vectorize an operation on a batched axis for one or more inputs. */
+declare const vmap: <F extends (...args: any[]) => JsTree<Array>>(f: WithArgsSubtype<F, JsTree<ArrayLike>>, inAxes?: number | MapJsTree<Parameters<F>, ArrayLike, number | null>) => (...args: MapJsTree<Parameters<F>, ArrayLike, ArrayLike>) => ReturnType<F>;
+/** Compute the Jacobian evaluated column-by-column by forward-mode AD. */
+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;
+};
+declare const jit: <F extends (...args: any[]) => JsTree<Array>>(f: WithArgsSubtype<F, JsTree<ArrayLike>>) => (...args: MapJsTree<Parameters<F>, ArrayLike, ArrayLike>) => ReturnType<F>;
+/**
+ * Produce a local linear approximation to a function at a point using jvp() and
+ * partial evaluation.
+ */
+declare const linearize: <F extends (...args: any[]) => JsTree<Array>>(f: WithArgsSubtype<F, JsTree<ArrayLike>>, ...primals: MapJsTree<Parameters<F>, ArrayLike, ArrayLike>) => [ReturnType<F>, (...tangents: MapJsTree<Parameters<F>, ArrayLike, ArrayLike>) => ReturnType<F>];
+/** Calculate the reverse-mode vector-Jacobian product for a function. */
+declare const vjp: <F extends (...args: any[]) => JsTree<Array>>(f: WithArgsSubtype<F, JsTree<ArrayLike>>, ...primals: MapJsTree<Parameters<F>, ArrayLike, ArrayLike>) => [ReturnType<F>, (cotangents: MapJsTree<ReturnType<F>, Array, ArrayLike>) => MapJsTree<Parameters<F>, ArrayLike, Array>];
+/**
+ * Compute the gradient of a scalar-valued function `f` with respect to its
+ * 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>;
+/** 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 };