license import { backend_util , GatherV2 , GatherV2Attrs , GatherV2Inputs , KernelConfig , scalar , Tensor , TypedArray , util } from '@tensorflow/tfjs' ; import { createTensorsTypeOpAttr , NodeJSKernelBackend } from '../nodejs_kernel_backend' ; export const gatherV2Config = { kernelName : GatherV2 , backendName : 'tensorflow' , kernelFunc : ( args ) => { const { x , indices } = args . inputs as GatherV2Inputs ; const backend = args . backend as NodeJSKernelBackend ; const { axis , batchDims } = args . attrs as { } as GatherV2Attrs ; const indicesVals = backend . readSync ( indices . dataId ) as TypedArray ; const axisDim = x . shape [ axis ] ; for ( let i = 0 ; i < indicesVals . length ; ++ i ) { const index = indicesVals [ i ] ; util . assert ( index <= axisDim - 1 && index >= 0 , ( ) => `GatherV2: the index value ${ index } is not in [0, ${ axisDim - 1 }]` ) ; } backend_util . segment_util . collectGatherOpShapeInfo ( x as Tensor , indices as Tensor , axis , batchDims ) ; const axisTensor = scalar ( axis , 'int32' ) ; const opAttrs = [ { name : 'batch_dims' , type : backend . binding . TF_ATTR_INT , value : batchDims } , createTensorsTypeOpAttr ( 'Tparams' , x . dtype ) , createTensorsTypeOpAttr ( 'Tindices' , indices . dtype ) , createTensorsTypeOpAttr ( 'Taxis' , 'int32' ) ] ; const res = backend . executeSingleOutput ( GatherV2 , opAttrs , [ x , indices , axisTensor ] ) ; axisTensor . dispose ( ) ; return res ; } } ;