Mojo function

get_start_and_end_for_partitions

get_start_and_end_for_partitions[tile_size: Int](num_keys: Int, num_partitions: Int, partition_idx: Int) -> Tuple[Int, Int]

Calculate start and end indices for a partition.

Non-empty partitions are packed at low indices 0..N-1 with partition_size = max(tile_size, align_up(ceildiv(num_keys, num_partitions), tile_size)); partitions >= N are empty (start == end == num_keys).

Args:

• ​num_keys (Int): Total number of keys (sequence length).
• ​num_partitions (Int): Number of partitions to split keys into.
• ​partition_idx (Int): Index of current partition (0 to num_partitions-1).

Returns:

Tuple[Int, Int]: Tuple of (start_idx, end_idx) for the partition, aligned to tile_size.

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit