license import { buffer , Rank , ShapeMap , TensorBuffer , TypedArray } from '@tensorflow/tfjs-core' ; export function scatterImpl < R extends Rank > ( indices , updates , shape , outputSize , sliceSize , numUpdates , sliceRank , strides , defaultValue , sumDupeIndices ) { const flattenShape = [ outputSize / sliceSize , sliceSize ] ; const indicesData = indices . values as TypedArray ; const updatesData = updates . values as TypedArray ; if ( outputSize === 0 ) { return buffer ( shape as ShapeMap [ R ] , updates . dtype ) ; } const outBuf = buffer ( flattenShape , updates . dtype ) ; ( outBuf . values as TypedArray ) . fill ( defaultValue ) ; for ( let i = 0 ; i < numUpdates ; i ++ ) { const index = [ ] ; let flattenIndex = 0 ; for ( let j = 0 ; j < sliceRank ; j ++ ) { const dim = indicesData [ i * sliceRank + j ] ; index . push ( dim ) ; flattenIndex += dim * strides [ j ] ; } if ( flattenIndex < 0 || flattenIndex >= outputSize / sliceSize ) { throw new Error ( `Invalid indices: ${ index } does not index into ${ shape }` ) ; } for ( let k = 0 ; k < sliceSize ; k ++ ) { if ( sumDupeIndices ) { outBuf . values [ flattenIndex * sliceSize + k ] += updatesData [ i * sliceSize + k ] ; } else { outBuf . values [ flattenIndex * sliceSize + k ] = updates . rank === 0 ? updatesData [ 0 ] : updatesData [ i * sliceSize + k ] ; } } } return outBuf as TensorBuffer < R > ; }