license import { backend_util , Conv2D , Conv2DAttrs , Conv2DInputs , env , KernelConfig , KernelFunc , TensorInfo } from '@tensorflow/tfjs-core' ; import { MathBackendWebGL } from '../backend_webgl' ; import { Conv2DProgram } from '../conv_gpu' ; import { conv2dByMatMul , conv2dWithIm2Row } from './Conv2D_impl' ; import { reshape } from './Reshape' ; export function conv2d ( args ) { const { inputs , backend , attrs } = args ; const { x , filter } = inputs ; const { strides , pad , dataFormat , dilations , dimRoundingMode } = attrs ; const $dataFormat = backend_util . convertConv2DDataFormat ( dataFormat ) ; const convInfo = backend_util . computeConv2DInfo ( x . shape as [ number , number , number , number ] , filter . shape as [ number , number , number , number ] , strides , dilations , pad , dimRoundingMode , false , $dataFormat ) ; let out ; if ( convInfo . filterHeight === 1 && convInfo . filterWidth === 1 && convInfo . dilationHeight === 1 && convInfo . dilationWidth === 1 && convInfo . strideHeight === 1 && convInfo . strideWidth === 1 && ( convInfo . padInfo . type === 'SAME' || convInfo . padInfo . type === 'VALID' ) ) { out = conv2dByMatMul ( { x , filter , convInfo , backend } ) ; } else if ( env ( ) . getBool ( 'WEBGL_CONV_IM2COL' ) && x . shape [ 0 ] === 1 ) { out = conv2dWithIm2Row ( { x , filter , convInfo , backend } ) ; } else { const program = new Conv2DProgram ( convInfo ) ; out = backend . runWebGLProgram ( program , [ x , filter ] , 'float32' ) ; } const outReshaped = reshape ( { inputs : { x : out } , backend , attrs : { shape : convInfo . outShape } } ) ; backend . disposeIntermediateTensorInfo ( out ) ; return outReshaped ; } export const conv2DConfig = { kernelName : Conv2D , backendName : 'webgl' , kernelFunc : conv2d as { } as KernelFunc , } ;