license import { ENGINE } from '../engine' ; import { AvgPool3D , AvgPool3DAttrs , AvgPool3DInputs } from '../kernel_names' ; import { NamedAttrMap } from '../kernel_registry' ; import { Tensor4D , Tensor5D } from '../tensor' ; import { NamedTensorMap } from '../tensor_types' ; import { convertToTensor } from '../tensor_util_env' ; import { TensorLike } from '../types' ; import * as util from '../util' ; import { checkPadOnDimRoundingMode } from './conv_util' ; import { cast } from './cast' ; import { op } from './operation' ; import { reshape } from './reshape' ; param x param filterSize param strides param pad param dimRoundingMode param dataFormat doc function avgPool3d_ < T extends Tensor4D | Tensor5D > ( x , filterSize , strides , pad , dimRoundingMode ? , dataFormat = 'NDHWC' ) { const $x = convertToTensor ( x , 'x' , 'avgPool3d' , 'float32' ) ; let x5D = $x as Tensor5D ; let reshapedTo5D = false ; if ( $x . rank === 4 ) { reshapedTo5D = true ; x5D = reshape ( $x , [ 1 , $x . shape [ 0 ] , $x . shape [ 1 ] , $x . shape [ 2 ] , $x . shape [ 3 ] ] ) ; } util . assert ( x5D . rank === 5 , ( ) => `Error in avgPool3d: x must be rank 5 but got rank ${ x5D . rank }.` ) ; util . assert ( dataFormat === 'NDHWC' , ( ) => `Error in avgPool3d: Only NDHWC is currently supported, ` + `but got dataFormat of ${ dataFormat }` ) ; checkPadOnDimRoundingMode ( 'avgPool3d' , pad , dimRoundingMode ) ; const inputs = { x : x5D } ; const attrs = { filterSize , strides , pad , dimRoundingMode , dataFormat } ; let res = ENGINE . runKernel ( AvgPool3D , inputs as { } as NamedTensorMap , attrs as { } as NamedAttrMap ) as T ; res = cast ( res , x5D . dtype ) ; if ( reshapedTo5D ) { return reshape ( res , [ res . shape [ 1 ] , res . shape [ 2 ] , res . shape [ 3 ] , res . shape [ 4 ] ] ) as T ; } return res ; } export const avgPool3d = op ( { avgPool3d_ } ) ;