license import { backend_util , DataType , TypedArray , util } from '@tensorflow/tfjs-core' ; export function sparseSegmentReductionImpl ( input , inputShape , inputDType , indices , segmentIds , isMean = false , defaultValue = 0 ) { const numIndices = indices . length ; const inputFlat = [ inputShape [ 0 ] , input . length / inputShape [ 0 ] ] ; const numCol = inputFlat [ 1 ] ; const lastSegmentIdPlusOne = numIndices > 0 ? segmentIds [ numIndices - 1 ] + 1 : 0 ; const outputRows = lastSegmentIdPlusOne ; if ( outputRows < 0 ) { throw new Error ( backend_util . getSparseSegmentReductionNegativeSegmentIdsErrorMessage ( ) ) ; } const outputShape = inputShape . slice ( ) ; outputShape [ 0 ] = outputRows ; const outputLength = outputShape . reduce ( ( product , value ) => product * value , 1 ) ; const output = util . getArrayFromDType ( inputDType , outputLength ) as TypedArray ; if ( numIndices === 0 ) { if ( outputRows > 0 ) { output . fill ( defaultValue ) ; } return [ output , outputShape ] ; } if ( outputRows <= 0 ) { throw new Error ( backend_util . getSparseSegmentReductionNegativeSegmentIdsErrorMessage ( ) ) ; } let start = 0 , end = 1 ; let uninitializedIndex = 0 ; let outIndex = segmentIds [ start ] ; while ( true ) { let nextIndex = 0 ; if ( end < numIndices ) { nextIndex = segmentIds [ end ] ; if ( outIndex === nextIndex ) { ++ end ; continue ; } if ( outIndex >= nextIndex ) { throw new Error ( backend_util . getSparseSegmentReductionNonIncreasingSegmentIdsErrorMessage ( ) ) ; } } if ( outIndex < 0 || outIndex >= outputRows ) { throw new Error ( backend_util . getSparseSegmentReductionSegmentIdOutOfRangeErrorMessage ( outIndex , outputRows ) ) ; } if ( outIndex > uninitializedIndex ) { output . fill ( defaultValue , uninitializedIndex * numCol , outIndex * numCol ) ; } for ( let i = start ; i < end ; ++ i ) { const index = indices [ i ] ; if ( index < 0 || index >= inputFlat [ 0 ] ) { throw new Error ( backend_util . getSparseSegmentReductionIndicesOutOfRangeErrorMessage ( i , indices [ i ] , inputFlat [ 0 ] ) ) ; } for ( let j = 0 ; j < numCol ; j ++ ) { output [ outIndex * numCol + j ] += input [ index * numCol + j ] ; } } if ( isMean ) { for ( let j = 0 ; j < numCol ; j ++ ) { output [ outIndex * numCol + j ] /= end - start ; } } start = end ; ++ end ; uninitializedIndex = outIndex + 1 ; outIndex = nextIndex ; if ( end > numIndices ) { break ; } } if ( uninitializedIndex < outputRows ) { output . fill ( defaultValue , uninitializedIndex * numCol , outputRows * numCol ) ; } return [ output , outputShape ] ; }