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.
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 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!