license import { MaxPoolWithArgmax , MaxPoolWithArgmaxAttrs , MaxPoolWithArgmaxInputs } from '@tensorflow/tfjs-core' ; import { backend_util , KernelConfig , util } from '@tensorflow/tfjs-core' ; import { MathBackendWebGL } from '../backend_webgl' ; import { maxPoolWithArgmaxImpl } from './MaxPoolWithArgmax_impl' ; export const maxPoolWithArgmaxConfig = { kernelName : MaxPoolWithArgmax , backendName : 'webgl' , kernelFunc : ( { inputs , attrs , backend } ) => { const { x } = inputs as MaxPoolWithArgmaxInputs ; const { filterSize , strides , pad , includeBatchInIndex } = attrs as { } as MaxPoolWithArgmaxAttrs ; const webglBackend = backend as MathBackendWebGL ; util . assert ( x . shape . length === 4 , ( ) => `Error in maxPool: input must be rank 4 but got rank ${ x . shape . length }.` ) ; const dilations = [ 1 , 1 ] ; util . assert ( backend_util . eitherStridesOrDilationsAreOne ( strides , dilations ) , ( ) => 'Error in maxPool: 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 ) ; const [ result , indexes ] = maxPoolWithArgmaxImpl ( x , includeBatchInIndex , convInfo , webglBackend ) ; return [ result , indexes ] ; } } ;