license import { env } from './environment' ; import { getGlobal } from './global_util' ; import * as log from './log' ; import { NamedGradientMap } from './tape' ; import { Tensor } from './tensor' ; import { DataType , RecursiveArray } from './types' ; const kernelRegistry = getGlobal ( 'kernelRegistry' , ( ) => new Map < string , KernelConfig > ( ) ) ; const gradRegistry = getGlobal ( 'gradRegistry' , ( ) => new Map < string , GradConfig > ( ) ) ; export type DataId = object ; type AttributeValue = number | number [ ] | boolean | boolean [ ] | string | string [ ] | NamedAttrMap ; /** These are extra non-tensor/primitive params passed to kernel functions. */ export type Attribute = AttributeValue | RecursiveArray < AttributeValue > ; /** Specifies the code to run when executing a kernel. */ export type KernelFunc = ( params : { inputs : NamedTensorInfoMap ; backend : { } ; attrs ? : NamedAttrMap ; } ) => TensorInfo | TensorInfo [ ] ; /** The function to run when computing a gradient during backprop. */ export type GradFunc = ( dy : Tensor | Tensor [ ] , saved : Tensor [ ] , attrs : NamedAttrMap ) => NamedGradientMap ; /** Function that gets called after the backend initializes. */ export type KernelSetupFunc = ( backend : { } ) => void ; /** Function that gets called right before the backend is disposed. */ export type KernelDisposeFunc = KernelSetupFunc ; /** Config object for registering a kernel in the global registry. */ export interface KernelConfig { kernelName : string ; backendName : string ; kernelFunc : KernelFunc ; setupFunc ? : KernelSetupFunc ; disposeFunc ? : KernelDisposeFunc ; } /** Config object for registering a gradient in the global registry. */ export interface GradConfig { kernelName : string ; inputsToSave ? : string [ ] ; saveAllInputs ? : boolean ; outputsToSave ? : boolean [ ] ; gradFunc : GradFunc ; } /** Holds metadata for a given tensor. */ export interface TensorInfo { dataId : DataId ; shape : number [ ] ; dtype : DataType ; } export interface NamedTensorInfoMap { [ name ] : TensorInfo | undefined ; } export interface NamedAttrMap { [ name ] : Attribute ; } param kernelName param backendName export function getKernel ( kernelName , backendName ) { const key = makeKey ( kernelName , backendName ) ; return kernelRegistry . get ( key ) ; } param kernelName export function getGradient ( kernelName ) { return gradRegistry . get ( kernelName ) ; } export function getKernelsForBackend ( backendName ) { const it = kernelRegistry . entries ( ) ; const result = [ ] ; while ( true ) { const { done , value } = it . next ( ) ; if ( done ) { break ; } const [ key , config ] = value ; const [ backend , ] = key . split ( '_' ) ; if ( backend === backendName ) { result . push ( config ) ; } } return result ; } param config export function registerKernel ( config ) { const { kernelName , backendName } = config ; const key = makeKey ( kernelName , backendName ) ; if ( kernelRegistry . has ( key ) ) { log . warn ( `The kernel '${ kernelName }' for backend ` + `'${ backendName }' is already registered` ) ; } kernelRegistry . set ( key , config ) ; } param config export function registerGradient ( config ) { const { kernelName } = config ; if ( gradRegistry . has ( kernelName ) ) { if ( env ( ) . getBool ( 'DEBUG' ) ) { log . warn ( `Overriding the gradient for '${ kernelName }'` ) ; } } gradRegistry . set ( kernelName , config ) ; } param kernelName param backendName export function unregisterKernel ( kernelName , backendName ) { const key = makeKey ( kernelName , backendName ) ; if ( ! kernelRegistry . has ( key ) ) { throw new Error ( `The kernel '${ kernelName }' for backend ` + `'${ backendName }' is not registered` ) ; } kernelRegistry . delete ( key ) ; } /** Removes the registered gradient from the global registry. */ export function unregisterGradient ( kernelName ) { if ( ! gradRegistry . has ( kernelName ) ) { throw new Error ( `The gradient '${ kernelName }' for backend is not registered` ) ; } gradRegistry . delete ( kernelName ) ; } param registeredBackendName param newBackendName export function copyRegisteredKernels ( registeredBackendName , newBackendName ) { const kernels = getKernelsForBackend ( registeredBackendName ) ; kernels . forEach ( kernelConfig => { const newKernelConfig = Object . assign ( { } , kernelConfig , { backendName : newBackendName } ) ; registerKernel ( newKernelConfig ) ; } ) ; } function makeKey ( kernelName , backendName ) { return `${ backendName }_${ kernelName }` ; }