license import { AvgPool , AvgPoolAttrs , AvgPoolInputs , backend_util , KernelConfig , KernelFunc , TensorInfo , TypedArray , util } from '@tensorflow/tfjs-core' ; import { MathBackendCPU } from '../backend_cpu' ; import { assertNotComplex } from '../cpu_util' ; import { pool } from '../utils/pool_utils' ; import { identity } from './Identity' ; export function avgPool ( args ) { const { inputs , backend , attrs } = args ; const { x } = inputs ; assertNotComplex ( x , 'avgPool' ) ; const { filterSize , strides , pad , dimRoundingMode } = attrs ; const dilations = 1 ; util . assert ( backend_util . eitherStridesOrDilationsAreOne ( strides , dilations ) , ( ) => 'Error in avgPool: Either strides or dilations must be 1. ' + `Got strides ${ strides } and dilations '${ dilations }'` ) ; const convInfo = backend_util . computePool2DInfo ( x . shape as [ number , number , number , number ] , filterSize , strides , dilations , pad , dimRoundingMode ) ; let res ; if ( convInfo . filterWidth === 1 && convInfo . filterHeight === 1 && util . arraysEqual ( convInfo . inShape , convInfo . outShape ) ) { res = identity ( { inputs : { x } , backend } ) ; } else { const xValues = backend . data . get ( x . dataId ) . values as TypedArray ; const strides = util . computeStrides ( x . shape ) ; const buffer = pool ( xValues , x . shape , x . dtype , strides , convInfo , 'avg' ) ; res = backend . makeTensorInfo ( convInfo . outShape , x . dtype , buffer . values as TypedArray ) ; } return res ; } export const avgPoolConfig = { kernelName : AvgPool , backendName : 'cpu' , kernelFunc : avgPool as { } as KernelFunc } ;