license import { TensorInfo } from '../kernel_registry' ; import { Tensor } from '../tensor' ; import { computeStrides , sizeFromShape } from '../util' ; param x export function validateUpdateShape ( shape , indices , updates ) { const sliceDim = ( indices . rank > 1 ) ? indices . shape [ indices . rank - 1 ] : 1 ; const batchDim = ( indices . rank > 1 ) ? indices . rank - 1 : 1 ; const shapeError = 'Must have updates.shape = indices.shape[:batchDim] + ' + `shape[sliceDim:], got updates.shape: ${ updates . shape }` + `, indices.shape: ${ indices . shape }, shape: ${ shape }` + `, sliceDim: ${ sliceDim }, and batchDim: ${ batchDim }.` ; if ( updates . rank < batchDim ) { throw new Error ( shapeError + ` update.rank < ${ batchDim }. ` ) ; } if ( shape . length < sliceDim + ( updates . rank - batchDim ) ) { throw new Error ( shapeError + ` Output shape length < ${ sliceDim + ( updates . rank - batchDim ) }` ) ; } if ( updates . rank !== batchDim + shape . length - sliceDim ) { throw new Error ( shapeError + ` update.rank != ${ batchDim + shape . length - sliceDim }` ) ; } for ( let d = 0 ; d < batchDim ; ++ d ) { if ( updates . shape [ d ] !== indices . shape [ d ] ) { throw new Error ( shapeError + ` updates.shape[${ d }] (${ updates . shape [ d ] }) != indices.shape[${ d }] (${ indices . shape [ d ] }).` ) ; } } for ( let d = 0 ; d < updates . rank - batchDim ; ++ d ) { if ( updates . shape [ d + batchDim ] !== shape [ d + sliceDim ] ) { throw new Error ( shapeError + ` updates.shape[${ d + batchDim }] (${ updates . shape [ d + batchDim ] }) != shape[${ d + batchDim }] (${ shape [ d + batchDim ] })` ) ; } } } export interface ScatterShapeInfo { sliceRank : number ; numUpdates : number ; sliceSize : number ; strides : number [ ] ; outputSize : number ; } param update param indices param shape export function validateInput ( updates , indices , shape ) { if ( indices . rank < 1 ) { throw new Error ( 'tf.scatterND() expects the indices to be rank 1 or higher,' + ` but the rank was ${ indices . rank }.` ) ; } if ( updates . rank < 1 ) { throw new Error ( 'tf.scatterND() expects the updates to be rank 1 or higher,' + ` but the rank was ${ updates . rank }.` ) ; } if ( indices . dtype !== 'int32' ) { throw new Error ( `The dtype of 'indices' should be int32, but got dtype: ${ indices . dtype }` ) ; } if ( shape . length < 1 ) { throw new Error ( `Output rank must be greater or equal to 1, but got shape: ${ shape }` ) ; } if ( shape . length === 0 ) { if ( indices . size === 0 ) { throw new Error ( `Indices specified for empty output. indices shape: ${ indices . shape }` ) ; } if ( updates . size === 0 ) { throw new Error ( `Updates specified for empty output. updates shape: ${ updates . shape }` ) ; } } validateUpdateShape ( shape , indices , updates ) ; } param update param indices param shape returns export function calculateShapes ( updates , indices , shape ) { const indicesRank = indices . shape . length ; const sliceRank = ( indicesRank > 1 ) ? indices . shape [ indicesRank - 1 ] : 1 ; const totalNd = shape . length ; let sliceSize = 1 ; for ( let i = sliceRank ; i < totalNd ; ++ i ) { sliceSize *= shape [ i ] ; } const safeSliceDim = ( sliceRank < 1 ) ? 1 : sliceRank ; const numUpdates = sizeFromShape ( indices . shape ) / safeSliceDim ; const strides = [ ... computeStrides ( shape . slice ( 0 , sliceRank ) ) , 1 ] ; const outputSize = sizeFromShape ( shape ) ; return { sliceRank , numUpdates , sliceSize , strides , outputSize } ; }