import { Buffer } from '../device.ts';
import type { DType } from '../dtype.ts';
import { ArrayMap, DefaultMap, WeakKeyMap } from '../helpers/helpers.ts';
import { UOp, type Variable } from '../ops.ts';
import type { ShapeTracker } from '../shape/shapetracker.ts';
import { ExecItem, Runner } from './realize.ts';
export declare class GraphException extends Error {
}
export declare const apply_graph_to_jit: (jit_cache: ExecItem[], input_rawbuffers: Buffer[], var_vals: Map<Variable, number>, max_batch_size?: number) => ExecItem[];
export declare const get_input_replace: (jit_cache: ExecItem[], input_rawbuffers: Buffer[]) => ArrayMap<[number, number], number>;
export declare class GraphRunner extends Runner {
    jit_cache: ExecItem[];
    input_replace: ArrayMap<[number, number], number>;
    var_vals_replace: Map<number, number[]>;
    launch_dims_replace: Map<number, [(number | undefined)?, (number | undefined)?]>;
    launch_dims_base: Map<number, [number[], number[]]>;
    vars: UOp[];
    symbolic_dims: ArrayMap<(number[] | undefined), number>;
    w_dependency_map: Map<bigint, any>;
    r_dependency_map: DefaultMap<bigint, any[]>;
    constructor(jit_cache: ExecItem[], input_rawbuffers: Buffer[], var_vals: Map<Variable, number>);
    updated_vars(var_vals: Map<Variable, number>): Generator<(number | undefined)[], void, unknown>;
    updated_launch_dims(var_vals: Map<Variable, number>): Generator<(number | number[])[], void, unknown>;
    _access_resources: (rawbufs: Buffer[], write: number[], new_dependency: any) => never;
}
export declare class MultiGraphRunner extends GraphRunner {
}
export declare const update_depends: (depends: Set<Buffer | undefined>, jit_cache: ExecItem[]) => void;
export declare class CapturedJit<Return extends any> {
    ret: Return;
    jit_cache: ExecItem[];
    input_replace: ArrayMap<[number, number], number>;
    extra_view_inputs: [number, number, string, number, DType][];
    expected_names: number[];
    expected_st_vars_dtype_device: [ShapeTracker, Variable[], DType, string | string[]][];
    constructor(ret: Return, // includes the Tensors or any other returned object
    jit_cache: ExecItem[], input_replace: ArrayMap<[number, number], number>, extra_view_inputs: [number, number, string, number, DType][], expected_names: number[], expected_st_vars_dtype_device: [ShapeTracker, Variable[], DType, string | string[]][]);
    _jit_cache: ExecItem[];
    _input_replace: ArrayMap<[number, number], number>;
    _first_run: boolean;
    init: () => void;
    _clear_inputs: () => void;
    free_intermediates: () => void;
    call: (input_buffers: Buffer[], var_vals: Map<Variable, number>) => Promise<Return>;
}
export declare const _prepare_jit_inputs: (...args: any[]) => Promise<readonly [Buffer<object>[], Map<UOp, number>, number[], [ShapeTracker, UOp[], DType, string | string[]][]]>;
export declare class TinyJit<Args extends any[], Return extends any> {
    fxn?: ((...a: Args) => Return | Promise<Return>) | undefined;
    captured?: CapturedJit<Return> | undefined;
    prune: boolean;
    cnt: number;
    _buffer_replace?: WeakKeyMap<Buffer, Buffer>;
    _jit_cache?: ExecItem[];
    constructor(fxn?: ((...a: Args) => Return | Promise<Return>) | undefined, captured?: CapturedJit<Return> | undefined, prune?: boolean);
    add_buffer: (b: Buffer) => Buffer;
    add: (ei: ExecItem) => void;
    reset: () => void;
    get jit_cache(): ExecItem[];
    get input_replace(): ArrayMap<[number, number], number>;
    call: (...args: Args) => Promise<Return>;
}