Mojo function

fold

fold[dtype: DType, input_dim: DimList, output_dim: DimList, //, stride: Tuple[Int, Int], dilation: Tuple[Int, Int], padding: Tuple[Int, Int], target: StringSlice[StaticConstantOrigin]](input: NDBuffer[dtype, 3, MutableAnyOrigin, input_dim], output: NDBuffer[dtype, 4, MutableAnyOrigin, output_dim], 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.
input_dim (DimList): The static shape of the input NDBuffer.
output_dim (DimList): The static shape of the output NDBuffer.
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 (NDBuffer[dtype, 3, MutableAnyOrigin, input_dim]): Input tensor to fold, shape [N, C x kernel size, num_blocks].
output (NDBuffer[dtype, 4, MutableAnyOrigin, output_dim]): 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.