license import { ENGINE } from '../engine' ; import { AvgPool , AvgPoolAttrs , AvgPoolInputs } from '../kernel_names' ; import { NamedAttrMap } from '../kernel_registry' ; import { Tensor3D , Tensor4D } from '../tensor' ; import { NamedTensorMap } from '../tensor_types' ; import { convertToTensor } from '../tensor_util_env' ; import { TensorLike } from '../types' ; import * as util from '../util' ; import { cast } from './cast' ; import * as conv_util from './conv_util' ; import { op } from './operation' ; import { reshape } from './reshape' ; param x param filterSize param strides param pad param dimRoundingMode function avgPool_ < T extends Tensor3D | Tensor4D > ( x , filterSize , strides , pad , dimRoundingMode ? ) { const $x = convertToTensor ( x , 'x' , 'avgPool' , 'float32' ) ; const dilations = 1 ; util . assert ( conv_util . eitherStridesOrDilationsAreOne ( strides , dilations ) , ( ) => 'Error in avgPool: Either strides or dilations must be 1. ' + `Got strides ${ strides } and dilations '${ dilations }'` ) ; let x4D = $x as Tensor4D ; let reshapedTo4D = false ; if ( $x . rank === 3 ) { reshapedTo4D = true ; x4D = reshape ( $x , [ 1 , $x . shape [ 0 ] , $x . shape [ 1 ] , $x . shape [ 2 ] ] ) ; } util . assert ( x4D . rank === 4 , ( ) => `Error in avgPool: x must be rank 4 but got rank ${ x4D . rank }.` ) ; conv_util . checkPadOnDimRoundingMode ( 'avgPool' , pad , dimRoundingMode ) ; const inputs = { x : x4D } ; const attrs = { filterSize , strides , pad , dimRoundingMode } ; let res = ENGINE . runKernel ( AvgPool , inputs as { } as NamedTensorMap , attrs as { } as NamedAttrMap ) as T ; res = cast ( res , $x . dtype ) ; if ( reshapedTo4D ) { return reshape ( res , [ res . shape [ 1 ] , res . shape [ 2 ] , res . shape [ 3 ] ] ) as T ; } return res ; } export const avgPool = op ( { avgPool_ } ) ;