license import { ENGINE } from '../engine' ; import { Conv3D , Conv3DAttrs , Conv3DInputs } 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 { eitherStridesOrDilationsAreOne } from './conv_util' ; import { op } from './operation' ; import { reshape } from './reshape' ; param x param filter param strides param pad param dataFormat param dilations doc function conv3d_ < T extends Tensor4D | Tensor5D > ( x , filter , strides , pad , dataFormat = 'NDHWC' , dilations = [ 1 , 1 , 1 ] ) { const $x = convertToTensor ( x , 'x' , 'conv3d' ) ; const $filter = convertToTensor ( filter , 'filter' , 'conv3d' ) ; 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 conv3d: input must be rank 5, but got rank ${ x5D . rank }.` ) ; util . assert ( $filter . rank === 5 , ( ) => `Error in conv3d: filter must be rank 5, but got rank ` + `${ $filter . rank }.` ) ; util . assert ( x5D . shape [ 4 ] === $filter . shape [ 3 ] , ( ) => `Error in conv3d: depth of input (${ x5D . shape [ 4 ] }) must match ` + `input depth for filter ${ $filter . shape [ 3 ] }.` ) ; util . assert ( eitherStridesOrDilationsAreOne ( strides , dilations ) , ( ) => 'Error in conv3D: Either strides or dilations must be 1. ' + `Got strides ${ strides } and dilations '${ dilations }'` ) ; util . assert ( dataFormat === 'NDHWC' , ( ) => `Error in conv3d: got dataFormat of ${ dataFormat } but only NDHWC is currently supported.` ) ; const inputs = { x : x5D , filter : $filter } ; const attrs = { strides , pad , dataFormat , dilations } ; const res = ENGINE . runKernel ( Conv3D , inputs as { } as NamedTensorMap , attrs as { } as NamedAttrMap ) as T ; if ( reshapedTo5D ) { return reshape ( res , [ res . shape [ 1 ] , res . shape [ 2 ] , res . shape [ 3 ] , res . shape [ 4 ] ] ) as T ; } return res ; } export const conv3d = op ( { conv3d_ } ) ;