license import { backend_util , DataType , KernelConfig , KernelFunc , Prod , ProdAttrs , ProdInputs , TensorInfo , TypedArray , upcastType , util } from '@tensorflow/tfjs-core' ; import { MathBackendCPU } from '../backend_cpu' ; import { assertNotComplex } from '../cpu_util' ; import { transpose } from './Transpose' ; export function prodImpl ( xShape , xDtype , xVals , reductionAxes ) { const [ outShape , reduceShape ] = backend_util . computeOutAndReduceShapes ( xShape , reductionAxes ) ; const outDtype = upcastType ( xDtype , 'int32' ) ; const outVals = util . makeZerosTypedArray ( util . sizeFromShape ( outShape ) , outDtype ) as TypedArray ; const reduceSize = util . sizeFromShape ( reduceShape ) ; for ( let i = 0 ; i < outVals . length ; ++ i ) { const offset = i * reduceSize ; let prod = 1 ; for ( let j = 0 ; j < reduceSize ; ++ j ) { prod *= xVals [ offset + j ] ; } outVals [ i ] = prod ; } return { outVals , outShape , outDtype } ; } export function prod ( args ) { const { inputs , backend , attrs } = args ; const { x } = inputs ; const { axis , keepDims } = attrs ; assertNotComplex ( x , 'prod' ) ; const xRank = x . shape . length ; const axes = util . parseAxisParam ( axis , x . shape ) ; const permutation = backend_util . getAxesPermutation ( axes , xRank ) ; let reductionAxes = axes ; let permutedX = x ; const intermediateTensorInfos = [ ] ; if ( permutation != null ) { permutedX = transpose ( { inputs : { x } , backend , attrs : { perm : permutation } } ) ; intermediateTensorInfos . push ( permutedX ) ; reductionAxes = backend_util . getInnerMostAxes ( reductionAxes . length , xRank ) ; } const xVals = backend . data . get ( permutedX . dataId ) . values as TypedArray ; const { outVals , outShape , outDtype } = prodImpl ( permutedX . shape , permutedX . dtype , xVals , reductionAxes ) ; let resultShape = outShape ; if ( keepDims ) { resultShape = backend_util . expandShapeToKeepDim ( outShape , axes ) ; } intermediateTensorInfos . forEach ( t => backend . disposeIntermediateTensorInfo ( t ) ) ; return backend . makeTensorInfo ( resultShape , outDtype , outVals ) ; } export const prodConfig = { kernelName : Prod , backendName : 'cpu' , kernelFunc : prod as { } as KernelFunc } ;