license import { KernelConfig , OneHot , OneHotAttrs , OneHotInputs , scalar } from '@tensorflow/tfjs' ; import { createTensorsTypeOpAttr , NodeJSKernelBackend } from '../nodejs_kernel_backend' ; export const oneHotConfig = { kernelName : OneHot , backendName : 'tensorflow' , kernelFunc : ( args ) => { const { indices } = args . inputs as OneHotInputs ; const backend = args . backend as NodeJSKernelBackend ; const { depth , onValue , offValue } = args . attrs as { } as OneHotAttrs ; const depthTensor = scalar ( depth , 'int32' ) ; const onValueTensor = scalar ( onValue , 'int32' ) ; const offValueTensor = scalar ( offValue , 'int32' ) ; const opAttrs = [ { name : 'axis' , type : backend . binding . TF_ATTR_INT , value : - 1 } , createTensorsTypeOpAttr ( 'T' , indices . dtype ) , createTensorsTypeOpAttr ( 'TI' , indices . dtype ) ] ; const res = backend . executeSingleOutput ( OneHot , opAttrs , [ indices , depthTensor , onValueTensor , offValueTensor ] ) ; depthTensor . dispose ( ) ; onValueTensor . dispose ( ) ; offValueTensor . dispose ( ) ; return res ; } } ;