For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo function

get_partition

def get_partition(task_id: Int, num_partitions: IndexList[Int(4)], conv_shape: ConvShape, micro_kernel_height: Int, micro_kernel_f_size: Int) -> ConvPartition

Computes the partition of the transposed convolution workload assigned to a single task.

Args:

• ​task_id (Int): Identifier of the task to compute the partition for.
• ​num_partitions (IndexList[Int(4)]): Number of partitions along each dimension.
• ​conv_shape (ConvShape): Shape descriptor of the transposed convolution.
• ​micro_kernel_height (Int): Height of the micro kernel in the output spatial dimension.
• ​micro_kernel_f_size (Int): Size of the micro kernel along the output channel dimension.

Returns:

ConvPartition: The partition describing the offset and size of the work assigned to the task.