Mojo function

non_max_suppression_shape_func

non_max_suppression_shape_func[dtype: DType](boxes: TileTensor[dtype, address_space=boxes.address_space, linear_idx_type=boxes.linear_idx_type, element_size=boxes.element_size], scores: TileTensor[dtype, address_space=scores.address_space, linear_idx_type=scores.linear_idx_type, element_size=scores.element_size], max_output_boxes_per_class: Int, iou_threshold: Float32, score_threshold: Float32) -> IndexList[2]

Compute the output shape for NMS without allocating the output buffer.

This function performs a dry-run of NMS to determine how many boxes will be selected, allowing proper output buffer allocation. Can be removed once the graph compiler supports value semantic kernels that allocate their own output.

Args:

boxes (TileTensor[dtype, address_space=boxes.address_space, linear_idx_type=boxes.linear_idx_type, element_size=boxes.element_size]): Rank-3 tensor of bounding boxes with shape (batch, num_boxes, 4).
scores (TileTensor[dtype, address_space=scores.address_space, linear_idx_type=scores.linear_idx_type, element_size=scores.element_size]): Rank-3 tensor of scores with shape (batch, num_classes, num_boxes).
max_output_boxes_per_class (Int): Maximum number of boxes to select per class.
iou_threshold (Float32): IoU threshold for suppression.
score_threshold (Float32): Minimum score threshold.

Returns:

IndexList[2]: A 2-element IndexList specifying the output shape (num_selected_boxes, 3).