UNPKG

@tensorflow/tfjs-core

Version:

Hardware-accelerated JavaScript library for machine intelligence

101 lines (100 loc) 3.8 kB
import { BackendTimingInfo, DataMover, KernelBackend } from './kernels/backend'; import { DataId, Tensor, Tensor3D, TensorTracker, Variable } from './tensor'; import { NamedTensorMap, NamedVariableMap, TensorContainer } from './tensor_types'; import { DataValues } from './types'; export declare type ForwardFunc<T> = (backend: KernelBackend, save?: <S extends Tensor>(tensor: S) => S) => T; export declare type CustomGradientFunc<T extends Tensor> = (...args: Tensor[]) => { value: T; gradFunc: (dy: T) => Tensor | Tensor[]; }; export declare type MemoryInfo = { numTensors: number; numDataBuffers: number; numBytes: number; unreliable?: boolean; reasons: string[]; }; declare type KernelProfile = { name: string; bytesAdded: number; totalBytesSnapshot: number; tensorsAdded: number; totalTensorsSnapshot: number; inputShapes: number[][]; outputShape: number[] | number[][]; }; export declare type ProfileInfo = { newBytes: number; newTensors: number; peakBytes: number; kernels: KernelProfile[]; result: TensorContainer; }; export interface TimingInfo extends BackendTimingInfo { wallMs: number; } export declare type ScopeFn<T extends TensorContainer> = () => T; export interface TensorManager { registerTensor(a: Tensor): void; registerVariable(v: Variable): void; disposeTensor(a: Tensor): void; memory(): { numDataBuffers: number; numBytes: number; }; } export declare class Engine implements TensorManager, TensorTracker, DataMover { backend: KernelBackend; safeMode: boolean; private debugMode; registeredVariables: NamedVariableMap; private nextTapeNodeId; private numBytes; private numTensors; private numStringTensors; private numDataBuffers; private profiling; private activeProfile; private activeTape; private gradientScopeCount; private customGradientDepth; private activeScope; private scopeStack; private keepTensors; private profiler; private tensorInfo; constructor(backend: KernelBackend, safeMode: boolean, debugMode: () => boolean); moveData(dataId: DataId): void; tidy<T extends TensorContainer>(nameOrFn: string | ScopeFn<T>, fn?: ScopeFn<T>, gradMode?: boolean): T; private scopedRun; private static nextTensorId; nextTensorId(): number; private static nextVariableId; nextVariableId(): number; runKernel<T extends Tensor | Tensor[], I extends NamedTensorMap>(forwardFunc: ForwardFunc<T>, inputs: I, backwardsFunc?: (dy: T, saved: Tensor[]) => { [P in keyof I]: () => I[P]; }): T; registerTensor(a: Tensor | Variable): void; registerVariable(v: Variable): void; disposeTensor(a: Tensor): void; disposeVariables(): void; memory(): MemoryInfo; profile(query: () => TensorContainer): Promise<ProfileInfo>; private shouldRecord; private addTapeNode; keep<T extends Tensor>(result: T): T; startScope(name?: string, gradientsMode?: boolean): void; endScope(result?: TensorContainer, gradientsMode?: boolean): void; gradients<T extends Tensor>(f: () => T, xs: Tensor[], dy?: T, allowNoGradients?: boolean): { value: T; grads: Tensor[]; }; customGrad<T extends Tensor>(f: CustomGradientFunc<T>): (...args: Tensor[]) => T; write(dataId: DataId, values: DataValues): void; readSync(dataId: DataId): DataValues; read(dataId: DataId): Promise<DataValues>; fromPixels(pixels: ImageData | HTMLImageElement | HTMLCanvasElement | HTMLVideoElement, numChannels: number): Tensor3D; time(query: () => void): Promise<TimingInfo>; private track; } export {};