license import { ENGINE } from '../engine' ; import { Conv2D , Conv2DAttrs , Conv2DInputs } 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 * as conv_util from './conv_util' ; import { op } from './operation' ; import { reshape } from './reshape' ; param x param filter param strides param pad param dataFormat param dilations param dimRoundingMode doc function conv2d_ < T extends Tensor3D | Tensor4D > ( x , filter , strides , pad , dataFormat = 'NHWC' , dilations = [ 1 , 1 ] , dimRoundingMode ? ) { const $x = convertToTensor ( x , 'x' , 'conv2d' , 'float32' ) ; const $filter = convertToTensor ( filter , 'filter' , 'conv2d' , 'float32' ) ; 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 conv2d: input must be rank 4, but got rank ${ x4D . rank }.` ) ; util . assert ( $filter . rank === 4 , ( ) => `Error in conv2d: filter must be rank 4, but got rank ` + `${ $filter . rank }.` ) ; conv_util . checkPadOnDimRoundingMode ( 'conv2d' , pad , dimRoundingMode ) ; const inDepth = dataFormat === 'NHWC' ? x4D . shape [ 3 ] : x4D . shape [ 1 ] ; util . assert ( inDepth === $filter . shape [ 2 ] , ( ) => `Error in conv2d: depth of input (${ inDepth }) must match ` + `input depth for filter ${ $filter . shape [ 2 ] }.` ) ; util . assert ( conv_util . eitherStridesOrDilationsAreOne ( strides , dilations ) , ( ) => 'Error in conv2D: Either strides or dilations must be 1. ' + `Got strides ${ strides } and dilations '${ dilations }'` ) ; const inputs = { x : x4D , filter : $filter } ; const attrs = { strides , pad , dataFormat , dilations , dimRoundingMode } ; const res = ENGINE . runKernel ( Conv2D , inputs as { } as NamedTensorMap , attrs as { } as NamedAttrMap ) as T ; if ( reshapedTo4D ) { return reshape ( res , [ res . shape [ 1 ] , res . shape [ 2 ] , res . shape [ 3 ] ] ) as T ; } return res ; } export const conv2d = op ( { conv2d_ } ) ;