license import { backend_util , buffer , KernelConfig , KernelFunc , MaxPoolGrad , MaxPoolGradAttrs , MaxPoolGradInputs , Rank , TensorInfo , TypedArray } from '@tensorflow/tfjs-core' ; import { MathBackendCPU } from '../backend_cpu' ; import { assertNotComplex } from '../cpu_util' ; import { maxPoolPositions } from '../utils/pool_utils' ; export function maxPoolGrad ( args ) { const { inputs , backend , attrs } = args ; const { dy , input , output } = inputs ; const x = input ; assertNotComplex ( [ input , output ] , 'maxPoolGrad' ) ; const { filterSize , strides , pad , dimRoundingMode } = attrs ; const convInfo = backend_util . computePool2DInfo ( x . shape as [ number , number , number , number ] , filterSize , strides , 1 , pad , dimRoundingMode ) ; const xValues = backend . data . get ( x . dataId ) . values as TypedArray ; const maxPosBuf = buffer ( convInfo . outShape , x . dtype , maxPoolPositions ( xValues , x . shape , x . dtype , convInfo ) . values ) ; const strideHeight = convInfo . strideHeight ; const strideWidth = convInfo . strideWidth ; const dilationHeight = convInfo . dilationHeight ; const dilationWidth = convInfo . dilationWidth ; const effectiveFilterHeight = convInfo . effectiveFilterHeight ; const effectiveFilterWidth = convInfo . effectiveFilterWidth ; const padLeft = effectiveFilterWidth - 1 - convInfo . padInfo . left ; const padTop = effectiveFilterHeight - 1 - convInfo . padInfo . top ; const dx = buffer < Rank . R4 > ( x . shape as [ number , number , number , number ] , 'float32' ) ; const dyData = backend . data . get ( dy . dataId ) . values as Float32Array ; const dyBuf = buffer < Rank . R4 > ( dy . shape as [ number , number , number , number ] , 'float32' , dyData ) ; for ( let b = 0 ; b < convInfo . batchSize ; ++ b ) { for ( let d = 0 ; d < convInfo . inChannels ; ++ d ) { for ( let dxR = 0 ; dxR < convInfo . inHeight ; ++ dxR ) { for ( let dxC = 0 ; dxC < convInfo . inWidth ; ++ dxC ) { const dyRCorner = dxR - padTop ; const dyCCorner = dxC - padLeft ; let dotProd = 0 ; for ( let wR = 0 ; wR < effectiveFilterHeight ; wR += dilationHeight ) { const dyR = ( dyRCorner + wR ) / strideHeight ; if ( dyR < 0 || dyR >= convInfo . outHeight || Math . floor ( dyR ) !== dyR ) { continue ; } for ( let wC = 0 ; wC < effectiveFilterWidth ; wC += dilationWidth ) { const dyC = ( dyCCorner + wC ) / strideWidth ; if ( dyC < 0 || dyC >= convInfo . outWidth || Math . floor ( dyC ) !== dyC ) { continue ; } const maxPos = effectiveFilterHeight * effectiveFilterWidth - 1 - ( maxPosBuf . get ( b , dyR , dyC , d ) as number ) ; const curPos = wR * effectiveFilterWidth + wC ; const mask = maxPos === curPos ? 1 : 0 ; if ( mask === 0 ) { continue ; } const pixel = dyBuf . get ( b , dyR , dyC , d ) ; dotProd += pixel * mask ; } } dx . set ( dotProd , b , dxR , dxC , d ) ; } } } } return backend . makeTensorInfo ( dx . shape , dx . dtype , dx . values ) ; } export const maxPoolGradConfig = { kernelName : MaxPoolGrad , backendName : 'cpu' , kernelFunc : maxPoolGrad as { } as KernelFunc } ;