license import { KernelConfig , scalar , UnsortedSegmentSum , UnsortedSegmentSumAttrs , UnsortedSegmentSumInputs } from '@tensorflow/tfjs' ; import { createTensorsTypeOpAttr , NodeJSKernelBackend } from '../nodejs_kernel_backend' ; export const unsortedSegmentSumConfig = { kernelName : UnsortedSegmentSum , backendName : 'tensorflow' , kernelFunc : ( args ) => { const { x , segmentIds } = args . inputs as UnsortedSegmentSumInputs ; const backend = args . backend as NodeJSKernelBackend ; const { numSegments } = args . attrs as { } as UnsortedSegmentSumAttrs ; const opAttrs = [ createTensorsTypeOpAttr ( 'T' , x . dtype ) , createTensorsTypeOpAttr ( 'Tindices' , 'int32' ) , createTensorsTypeOpAttr ( 'Tnumsegments' , 'int32' ) ] ; const numSegmentsT = scalar ( numSegments , 'int32' ) ; const res = backend . executeSingleOutput ( UnsortedSegmentSum , opAttrs , [ x , segmentIds , numSegmentsT ] ) ; numSegmentsT . dispose ( ) ; return res ; } } ;