license import { backend_util , DataType , TensorInfo } from '@tensorflow/tfjs-core' ; import { MathBackendWebGL } from '../backend_webgl' ; import { MeanProgram } from '../mean_gpu' ; import { ReduceProgram } from '../reduce_gpu' ; type ReduceTypes = 'all' | 'any' | 'max' | 'min' | 'sum' | 'prod' | 'mean' ; function getReductionStages ( inShape ) { const stages = [ ] ; while ( stages . length === 0 || stages [ stages . length - 1 ] . outSize !== 1 ) { const outSize = stages . length ? stages [ stages . length - 1 ] . outSize : inShape [ 1 ] ; const windowSize = backend_util . computeOptimalWindowSize ( outSize ) ; stages . push ( { inSize : outSize , windowSize , outSize : Math . ceil ( outSize / windowSize ) } ) ; } return stages ; } export function reduce ( x , dtype , reductionType , backend ) { const reductionStages = getReductionStages ( x . shape ) ; let result = x ; for ( let i = 0 ; i < reductionStages . length ; i ++ ) { const { inSize , windowSize , outSize } = reductionStages [ i ] ; let program ; let previousResult ; if ( reductionType === 'mean' ) { program = i === 0 ? new MeanProgram ( { windowSize , inSize , batchSize : x . shape [ 0 ] , outSize } , inSize ) : new MeanProgram ( { windowSize , inSize , batchSize : x . shape [ 0 ] , outSize } ) ; } else { program = new ReduceProgram ( { windowSize , inSize , batchSize : x . shape [ 0 ] , outSize } , reductionType ) ; } previousResult = result ; result = backend . runWebGLProgram ( program , [ result ] , dtype ) ; if ( previousResult . dataId !== x . dataId ) { backend . disposeIntermediateTensorInfo ( previousResult ) ; } } return result ; }