license import { backend_util , KernelConfig , KernelFunc , Mean , MeanAttrs , MeanInputs , TensorInfo , util } from '@tensorflow/tfjs-core' ; import { MathBackendCPU } from '../backend_cpu' ; import { cast } from './Cast' ; import { div } from './RealDiv' ; import { sum } from './Sum' ; export function mean ( args ) { const { inputs , backend , attrs } = args ; const { x } = inputs ; const { axis , keepDims } = attrs ; const axes = util . parseAxisParam ( axis , x . shape ) ; const shapes = backend_util . computeOutAndReduceShapes ( x . shape , axes ) ; const reduceShape = shapes [ 1 ] ; const reduceSize = util . sizeFromShape ( reduceShape ) ; const toDispose = [ ] ; const reduceSizeScalar = backend . makeTensorInfo ( [ ] , 'float32' , new Float32Array ( [ reduceSize ] ) ) ; toDispose . push ( reduceSizeScalar ) ; const $x = cast ( { inputs : { x } , backend , attrs : { dtype : 'float32' } } ) ; toDispose . push ( $x ) ; const res = div ( { inputs : { a : $x , b : reduceSizeScalar } , backend } ) as TensorInfo ; toDispose . push ( res ) ; const result = sum ( { inputs : { x : res } , backend , attrs : { axis , keepDims } } ) ; toDispose . forEach ( t => backend . disposeIntermediateTensorInfo ( t ) ) ; return result ; } export const meanConfig = { kernelName : Mean , backendName : 'cpu' , kernelFunc : mean as { } as KernelFunc } ;