conv2d_op

Structs

Struct: Conv2d

Namespace for 2D convolution operations.

Fields

Methods


compute_shape(mut curr: ArrayShape, args: List[ArrayShape])

Computes the shape of an array after a 2-dimensional convolution operation.

Args

curr: ArrayShape
args: List[ArrayShape]


__call__(mut curr: Array, args: List[Array])

more details

Args

curr: Array
args: List[Array]


vjp(primals: List[Array], grad: Array, out: Array) -> List[Array]

more details

Args

primals: List[Array]
grad: Array
out: Array

Returns

List[Array]


jvp(primals: List[Array], tangents: List[Array]) -> Array

more details

Args

primals: List[Array]
tangents: List[Array]

Returns

Array


fwd(arg0: Array, kernel: Array, bias: Array, stride: Tuple[Int, Int] = Tuple(VariadicPack(<store_to_mem({1}), store_to_mem({1})>, True)), padding: Tuple[Int, Int] = Tuple(VariadicPack(<store_to_mem({0}), store_to_mem({0})>, True)), dilation: Tuple[Int, Int] = Tuple(VariadicPack(<store_to_mem({1}), store_to_mem({1})>, True)), groups: Int = 1) -> Array

more details

Args

arg0: Array
kernel: Array
bias: Array
stride: Tuple[Int, Int] (default: Tuple(VariadicPack(<store_to_mem({1}), store_to_mem({1})>, True)))
padding: Tuple[Int, Int] (default: Tuple(VariadicPack(<store_to_mem({0}), store_to_mem({0})>, True)))
dilation: Tuple[Int, Int] (default: Tuple(VariadicPack(<store_to_mem({1}), store_to_mem({1})>, True)))
groups: Int (default: 1)

Returns

Array

Functions

conv2d


conv2d(arg0: Array, kernel: Array, bias: Array, stride: Tuple[Int, Int] = Tuple(VariadicPack(<store_to_mem({1}), store_to_mem({1})>, True)), padding: Tuple[Int, Int] = Tuple(VariadicPack(<store_to_mem({0}), store_to_mem({0})>, True)), dilation: Tuple[Int, Int] = Tuple(VariadicPack(<store_to_mem({1}), store_to_mem({1})>, True)), groups: Int = 1) -> Array

Applies a 2D convolution over an input image composed of several input planes.

Args

arg0: Array Input tensor of shape (batch_size, in_channels, height, width).
kernel: Array Convolution kernel of shape (out_channels, in_channels // groups, kernel_height, kernel_width).
bias: Array Bias tensor of shape (out_channels).
stride: Tuple[Int, Int] (default: Tuple(VariadicPack(<store_to_mem({1}), store_to_mem({1})>, True))) Stride of the convolution.
padding: Tuple[Int, Int] (default: Tuple(VariadicPack(<store_to_mem({0}), store_to_mem({0})>, True))) Zero-padding added to both sides of the input.
dilation: Tuple[Int, Int] (default: Tuple(VariadicPack(<store_to_mem({1}), store_to_mem({1})>, True))) Spacing between kernel elements.
groups: Int (default: 1) Number of blocked connections from input channels to output channels.

Returns

Array - Output tensor of shape (batch_size, out_channels, output_height, output_width).