Mojo function

fold

fold[dtype: DType, input_dim: DimList, output_dim: DimList, 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], stride: IndexList[2], dilation: IndexList[2], padding: IndexList[2], ctx: DeviceContextPtr)

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

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]): Spacial shape of the output tensor (H, W).
kernel_size (IndexList[2]): Size of the sliding blocks.
stride (IndexList[2]): Stride of the sliding blocks.
dilation (IndexList[2]): Dilation of the sliding blocks.
padding (IndexList[2]): 0-paddings to be added on both sides of the inputs.
ctx (DeviceContextPtr): DeviceContextPtr.