license import { KernelConfig , SparseSegmentSum , SparseSegmentSumInputs , TensorInfo , TypedArray } from '@tensorflow/tfjs-core' ; import { MathBackendCPU } from '../backend_cpu' ; import { sparseSegmentReductionImpl } from './SparseSegmentReduction_impl' ; 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 }` ) ; } if ( indices . shape [ 0 ] !== segmentIds . shape [ 0 ] ) { throw new Error ( `segmentIds and indices should have same size.` ) ; } const $data = backend . data . get ( data . dataId ) . values as TypedArray ; const $indices = backend . data . get ( indices . dataId ) . values as TypedArray ; const $segmentIds = backend . data . get ( segmentIds . dataId ) . values as TypedArray ; const [ outputData , outputDataShape ] = sparseSegmentReductionImpl ( $data , data . shape , data . dtype , $indices , $segmentIds ) ; return backend . makeTensorInfo ( outputDataShape , data . dtype , outputData ) ; } export const sparseSegmentSumConfig = { kernelName : SparseSegmentSum , backendName : 'cpu' , kernelFunc : sparseSegmentSum , } ;