license import { backend_util , DataType , TypedArray , util } from '@tensorflow/tfjs-core' ; export function sparseReshapeImpl ( inputIndices , inputIndicesShape , inputDType , inputShape , targetShape ) { const denseSize = util . sizeFromShape ( inputShape ) ; const nnz = inputIndicesShape [ 0 ] ; const outputRank = targetShape . length ; const outputShape = [ ] ; let product = 1 ; let unknownIndex = - 1 ; for ( let d = 0 ; d < outputRank ; ++ d ) { const size = targetShape [ d ] ; if ( size === - 1 ) { if ( unknownIndex !== - 1 ) { throw new Error ( backend_util . getSparseReshapeMultipleNegativeOneOutputDimErrorMessage ( unknownIndex , d ) ) ; } unknownIndex = d ; outputShape . push ( 1 ) ; } else { if ( size < 0 ) { throw new Error ( backend_util . getSparseReshapeNegativeOutputDimErrorMessage ( d , size ) ) ; } product *= size ; outputShape . push ( size ) ; } } if ( unknownIndex !== - 1 ) { if ( product <= 0 ) { throw new Error ( backend_util . getSparseReshapeEmptyTensorZeroOutputDimErrorMessage ( ) ) ; } const missing = Math . trunc ( denseSize / product ) ; if ( product * missing !== denseSize ) { throw new Error ( backend_util . getSparseReshapeInputOutputMultipleErrorMessage ( inputShape , outputShape ) ) ; } outputShape [ unknownIndex ] = missing ; } const outputSize = util . sizeFromShape ( outputShape ) ; if ( outputSize !== denseSize ) { throw new Error ( backend_util . getSparseReshapeInputOutputMismatchErrorMessage ( inputShape , outputShape ) ) ; } const inputRank = inputShape . length ; const inputStrides = [ ] ; if ( inputRank > 0 ) { inputStrides [ inputRank - 1 ] = 1 ; for ( let d = inputRank - 2 ; d >= 0 ; -- d ) { inputStrides [ d ] = inputStrides [ d + 1 ] * inputShape [ d + 1 ] ; } } const outputStrides = [ ] ; if ( outputRank > 0 ) { outputStrides [ outputRank - 1 ] = 1 ; for ( let d = outputRank - 2 ; d >= 0 ; -- d ) { outputStrides [ d ] = outputStrides [ d + 1 ] * outputShape [ d + 1 ] ; } } const newIndices = util . getArrayFromDType ( inputDType , nnz * outputRank ) as TypedArray ; for ( let i = 0 ; i < nnz ; ++ i ) { let id = 0 ; for ( let j = 0 ; j < inputRank ; ++ j ) { id += inputIndices [ i * inputRank + j ] * inputStrides [ j ] ; } for ( let j = 0 ; j < outputRank ; ++ j ) { newIndices [ i * outputRank + j ] = Math . trunc ( id / outputStrides [ j ] ) ; id %= outputStrides [ j ] ; } } return [ newIndices , [ nnz , outputRank ] , outputShape ] ; }