license import { KernelConfig , scalar , TopK , TopKAttrs , TopKInputs } from '@tensorflow/tfjs' ; import { isNullOrUndefined } from 'util' ; import { createTensorsTypeOpAttr , NodeJSKernelBackend } from '../nodejs_kernel_backend' ; export const topKConfig = { kernelName : TopK , backendName : 'tensorflow' , kernelFunc : ( args ) => { const { x } = args . inputs as TopKInputs ; const backend = args . backend as NodeJSKernelBackend ; const { k , sorted } = args . attrs as { } as TopKAttrs ; const kCount = isNullOrUndefined ( k ) ? 1 : k ; const isSorted = isNullOrUndefined ( sorted ) ? true : sorted ; const opAttrs = [ { name : 'sorted' , type : backend . binding . TF_ATTR_BOOL , value : isSorted } , createTensorsTypeOpAttr ( 'T' , x . dtype ) , ] ; const kTensor = scalar ( kCount , 'int32' ) ; const res = backend . executeMultipleOutputs ( 'TopKV2' , opAttrs , [ x , kTensor ] , 2 ) ; kTensor . dispose ( ) ; return res ; } } ;