license import { KernelConfig , KernelFunc , SparseSegmentSum , SparseSegmentSumInputs , TensorInfo , TypedArray } from '@tensorflow/tfjs-core' ; import { MathBackendWebGL } from '../backend_webgl' ; import { sparseSegmentReductionImplCPU } from '../kernel_utils/shared' ; export function sparseSegmentSum ( args ) { const { inputs , backend } = args ; const { data , indices , segmentIds } = inputs ; if ( data . shape . length < 1 ) { throw new Error ( `Data should be at least 1 dimensional but received scalar` ) ; } if ( indices . shape . length !== 1 ) { throw new Error ( `Indices should be a vector but received shape
             ${ indices . shape }` ) ; } if ( segmentIds . shape . length !== 1 ) { throw new Error ( `Segment ids should be a vector but received shape
             ${ segmentIds . shape }` ) ; } const $data = backend . readSync ( data . dataId ) as TypedArray ; const $indices = backend . readSync ( indices . dataId ) as TypedArray ; const $segmentIds = backend . readSync ( segmentIds . dataId ) as TypedArray ; const [ outputData , outputDataShape ] = sparseSegmentReductionImplCPU ( $data , data . shape , data . dtype , $indices , $segmentIds ) ; return backend . makeTensorInfo ( outputDataShape , data . dtype , outputData ) ; } export const sparseSegmentSumConfig = { kernelName : SparseSegmentSum , backendName : 'webgl' , kernelFunc : sparseSegmentSum as { } as KernelFunc , } ;