Mojo function

fold

fold[dtype: DType, stride: Tuple[Int, Int], dilation: Tuple[Int, Int], padding: Tuple[Int, Int], target: StringSlice[StaticConstantOrigin]](input: LayoutTensor[dtype, layout, origin, address_space=address_space, element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment], output: LayoutTensor[dtype, layout, origin, address_space=address_space, element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment], output_size: IndexList[2], kernel_size: IndexList[2], ctx: DeviceContextPtr)

Folds array of sliding local blocks into a single output tensor.

Parameters:

dtype (DType): The data type for the input and output.
stride (Tuple[Int, Int]): Stride of the sliding blocks.
dilation (Tuple[Int, Int]): Dilation of the sliding blocks.
padding (Tuple[Int, Int]): 0-paddings to be added on both sides of the inputs.
target (StringSlice[StaticConstantOrigin]): The target architecture to compile for.

Args:

input (LayoutTensor[dtype, layout, origin, address_space=address_space, element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment]): Input tensor to fold, shape [N, C x kernel size, num_blocks].
output (LayoutTensor[dtype, layout, origin, address_space=address_space, element_layout=element_layout, layout_int_type=layout_int_type, linear_idx_type=linear_idx_type, masked=masked, alignment=alignment]): Output tensor to write to, shape [N, C, H, W].
output_size (IndexList[2]): Spatial shape of the output tensor (H, W).
kernel_size (IndexList[2]): Size of the sliding blocks.
ctx (DeviceContextPtr): DeviceContextPtr.