⚡️ Speed up function nms by 12%

[codeflash-ai]

📄 12% (0.12x) speedup for nms in unstructured_inference/models/yolox.py

⏱️ Runtime : 57.5 milliseconds → 51.3 milliseconds (best of 61 runs)

📝 Explanation and details

The optimized code achieves a 12% speedup by eliminating redundant array indexing operations within the NMS loop.

Key Optimization:

Instead of repeatedly indexing order[1:] multiple times per iteration (5 times in the original code for coordinate comparisons plus area lookups), the optimized version extracts this slice once into a variable rest. This single extraction is then reused for all subsequent operations.

Why This Works:

1. Reduced Array Slicing Overhead: Each order[1:] operation creates a new array view and involves pointer arithmetic and bounds checking. By doing this once instead of 5+ times per iteration (across 2,237 loop iterations based on profiler data), we save significant overhead.

2. Improved Memory Access Pattern: The rest variable maintains better cache locality since we're reusing the same slice reference rather than recreating it multiple times.

3. Boolean Indexing vs np.where: The change from np.where(ovr <= nms_thr)[0] to direct boolean masking (mask = ovr <= nms_thr; order = rest[mask]) eliminates the function call overhead of np.where and the subsequent integer addition operation inds + 1. Line profiler shows this reduces time from ~14.6ms to ~10.5ms across the loop iterations.

Performance Impact:

Based on the line profiler results:

• The coordinate extraction lines (xx1, yy1, xx2, yy2) show modest improvements (~0.5-0.6ms total saved)
• The masking operation shows the biggest win: from 8.7ms (np.where) + 6.0ms (indexing) = 14.7ms down to 4.1ms + 3.6ms = 7.7ms—a ~7ms savings

Context & Impact:

The nms function is called from multiclass_nms_class_agnostic, which processes detection results after filtering by score threshold. Since this is in a post-processing path for object detection (YOLOX model), even small speedups compound when processing multiple images or video frames. The optimization is most beneficial for:

• Cases with many overlapping boxes (as shown by the 22-26% speedup in tests like test_one_threshold and test_boxes_with_zero_area)
• Large-scale scenarios with 100+ boxes (11-14% improvements consistently observed)

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 50 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Click to see Generated Regression Tests

import numpy as np  # used to build numeric inputs and inspect outputs
# imports
import pytest  # used for our unit tests
from numba import njit
from unstructured_inference.models.yolox import nms

def test_basic_two_non_overlapping():
    # Two clearly non-overlapping boxes; higher score should be first in result
    boxes = np.array([
        [0.0, 0.0, 10.0, 10.0],  # box 0
        [20.0, 20.0, 30.0, 30.0]  # box 1, no overlap with box 0
    ], dtype=np.float64)
    scores = np.array([0.5, 0.9], dtype=np.float64)  # box 1 has higher score
    codeflash_output = nms(boxes, scores, 0.5); kept = codeflash_output # 73.2μs -> 66.3μs (10.5% faster)

def test_single_box_returns_single_index():
    # Single box should always be kept
    boxes = np.array([[5.0, 5.0, 6.0, 6.0]], dtype=np.float64)
    scores = np.array([0.123], dtype=np.float64)
    codeflash_output = nms(boxes, scores, 0.5); kept = codeflash_output # 46.7μs -> 41.9μs (11.4% faster)

def test_empty_inputs_return_empty_list():
    # Zero boxes should return an empty Python list
    boxes = np.empty((0, 4), dtype=np.float64)
    scores = np.empty((0,), dtype=np.float64)
    codeflash_output = nms(boxes, scores, 0.5); kept = codeflash_output # 17.2μs -> 17.7μs (2.69% slower)

def test_identical_boxes_threshold_zero_suppresses_all_but_top():
    # Two identical boxes with different scores and threshold 0 should keep only top scoring
    boxes = np.array([
        [10.0, 10.0, 20.0, 20.0],
        [10.0, 10.0, 20.0, 20.0]
    ], dtype=np.float64)
    scores = np.array([0.8, 0.95], dtype=np.float64)  # index 1 is top
    codeflash_output = nms(boxes, scores, 0.0); kept = codeflash_output # 55.6μs -> 48.8μs (13.8% faster)

def test_identical_boxes_threshold_one_keeps_all():
    # With threshold 1.0 both identical boxes should be kept because overlap == 1 and <=1
    boxes = np.array([
        [1.0, 1.0, 4.0, 4.0],
        [1.0, 1.0, 4.0, 4.0]
    ], dtype=np.float64)
    scores = np.array([0.4, 0.7], dtype=np.float64)  # index 1 highest
    codeflash_output = nms(boxes, scores, 1.0); kept = codeflash_output # 83.0μs -> 68.0μs (22.1% faster)

def test_degenerate_zero_area_boxes_do_not_divide_by_zero_and_are_kept():
    # Create boxes with zero area (x2 < x1 and y2 < y1 produce widths/heights that when added +1 become zero)
    # For such boxes the area is 0 and denom can be zero, code should set overlap to 0 and keep both when threshold=0
    boxes = np.array([
        [2.0, 2.0, 1.0, 1.0],  # degenerate box A with area 0
        [2.0, 2.0, 1.0, 1.0]   # degenerate box B identical to A
    ], dtype=np.float64)
    scores = np.array([0.2, 0.3], dtype=np.float64)  # index 1 higher
    codeflash_output = nms(boxes, scores, 0.0); kept = codeflash_output # 82.2μs -> 76.1μs (8.02% faster)

def test_partial_overlap_clipping_to_zero_width_height():
    # Two boxes positioned so that intersection clips to negative widths/heights -> intersection becomes zero
    boxes = np.array([
        [0.0, 0.0, 10.0, 10.0],  # box 0
        [11.0, 11.0, 12.0, 12.0]  # box 1, just outside box 0 by 1 -> no overlap after clipping
    ], dtype=np.float64)
    scores = np.array([0.6, 0.7], dtype=np.float64)
    codeflash_output = nms(boxes, scores, 0.0); kept = codeflash_output # 82.5μs -> 66.2μs (24.5% faster)

def test_large_scale_properties_and_determinism():
    # Generate a moderately large set of boxes to test performance and basic properties (n < 1000)
    rng = np.random.RandomState(42)  # deterministic RNG for reproducibility
    n = 500  # within allowed limit
    # Generate random top-left coordinates and sizes but ensure x1 <= x2 and y1 <= y2
    x1 = rng.randint(0, 200, size=n).astype(np.float64)
    y1 = rng.randint(0, 200, size=n).astype(np.float64)
    widths = rng.randint(1, 20, size=n).astype(np.float64)
    heights = rng.randint(1, 20, size=n).astype(np.float64)
    boxes = np.stack([x1, y1, x1 + widths, y1 + heights], axis=1)
    # random scores in [0,1)
    scores = rng.rand(n).astype(np.float64)

    # Run nms twice to ensure deterministic output (no random state dependence)
    codeflash_output = nms(boxes, scores, 0.5); kept1 = codeflash_output # 13.0ms -> 11.7ms (11.7% faster)
    codeflash_output = nms(boxes, scores, 0.5); kept2 = codeflash_output # 13.0ms -> 11.7ms (11.6% faster)

    # Scores of kept indices must be non-increasing (selection order is by descending score)
    kept_scores = [scores[i] for i in kept1]
    # check monotonic non-increasing property
    for a, b in zip(kept_scores, kept_scores[1:]):
        pass

def test_scores_with_ties_stable_descending_order_by_algorithm():
    # If scores tie, argsort + reverse produces a predictable ordering.
    # Construct 4 boxes with same score but different indices; nms should still return all when threshold=1.0
    boxes = np.array([
        [0.0, 0.0, 2.0, 2.0],
        [1.0, 1.0, 3.0, 3.0],
        [2.0, 2.0, 4.0, 4.0],
        [3.0, 3.0, 5.0, 5.0]
    ], dtype=np.float64)
    scores = np.array([0.5, 0.5, 0.5, 0.5], dtype=np.float64)  # all identical
    codeflash_output = nms(boxes, scores, 1.0); kept = codeflash_output # 124μs -> 111μs (11.1% faster)

def test_input_non_contiguous_arrays_are_handled():
    # Create Fortran-contiguous arrays (transpose of C-contiguous) to test ascontiguouspath
    base_boxes = np.array([
        [0.0, 0.0, 4.0, 4.0],
        [1.0, 1.0, 5.0, 5.0],
        [2.0, 2.0, 6.0, 6.0]
    ], dtype=np.float64)
    # Make a non-contiguous view by transposing then transposing back in a way that might produce non-contiguity
    boxes = base_boxes.T.T  # simple way but keep dtype; ensure it's at least a view
    # Force non-contiguity by slicing step
    boxes_noncontig = boxes[:, ::1]  # may still be contiguous but passes through array handling
    scores = np.array([0.1, 0.9, 0.5], dtype=np.float64)
    codeflash_output = nms(boxes_noncontig, scores, 0.3); kept = codeflash_output # 52.6μs -> 48.0μs (9.60% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import numpy as np
import pytest
from unstructured_inference.models.yolox import nms

class TestNMSBasicFunctionality:
    """Test basic functionality of the NMS function with normal inputs."""

    def test_single_box(self):
        """Test NMS with a single bounding box."""
        boxes = np.array([[10.0, 10.0, 20.0, 20.0]], dtype=np.float64)
        scores = np.array([0.9], dtype=np.float64)
        codeflash_output = nms(boxes, scores, 0.5); result = codeflash_output # 46.6μs -> 40.8μs (14.4% faster)

    def test_two_non_overlapping_boxes(self):
        """Test NMS with two boxes that don't overlap."""
        boxes = np.array(
            [[10.0, 10.0, 20.0, 20.0], [50.0, 50.0, 60.0, 60.0]], dtype=np.float64
        )
        scores = np.array([0.9, 0.8], dtype=np.float64)
        codeflash_output = nms(boxes, scores, 0.5); result = codeflash_output # 72.2μs -> 76.9μs (6.13% slower)

    def test_two_overlapping_boxes_high_overlap(self):
        """Test NMS with two boxes that have high overlap (one suppressed)."""
        boxes = np.array(
            [[10.0, 10.0, 30.0, 30.0], [15.0, 15.0, 35.0, 35.0]], dtype=np.float64
        )
        scores = np.array([0.9, 0.8], dtype=np.float64)
        codeflash_output = nms(boxes, scores, 0.3); result = codeflash_output # 54.1μs -> 47.8μs (13.1% faster)

    def test_multiple_boxes_with_mixed_overlaps(self):
        """Test NMS with multiple boxes having varied overlap relationships."""
        boxes = np.array(
            [
                [10.0, 10.0, 20.0, 20.0],  # box 0
                [15.0, 15.0, 25.0, 25.0],  # box 1 (overlaps with 0)
                [50.0, 50.0, 60.0, 60.0],  # box 2 (no overlap)
                [12.0, 12.0, 22.0, 22.0],  # box 3 (overlaps with 0)
            ],
            dtype=np.float64,
        )
        scores = np.array([0.95, 0.8, 0.9, 0.7], dtype=np.float64)
        codeflash_output = nms(boxes, scores, 0.3); result = codeflash_output # 102μs -> 94.9μs (8.41% faster)

    def test_identical_boxes_different_scores(self):
        """Test NMS with identical boxes but different confidence scores."""
        boxes = np.array(
            [[10.0, 10.0, 20.0, 20.0], [10.0, 10.0, 20.0, 20.0]], dtype=np.float64
        )
        scores = np.array([0.9, 0.5], dtype=np.float64)
        codeflash_output = nms(boxes, scores, 0.5); result = codeflash_output # 54.8μs -> 48.5μs (13.0% faster)

    def test_nms_threshold_boundary(self):
        """Test NMS with overlap exactly at threshold boundary."""
        boxes = np.array(
            [[0.0, 0.0, 10.0, 10.0], [5.0, 5.0, 15.0, 15.0]], dtype=np.float64
        )
        scores = np.array([0.9, 0.8], dtype=np.float64)
        codeflash_output = nms(boxes, scores, 0.5); result = codeflash_output # 83.2μs -> 67.5μs (23.3% faster)

    def test_return_type_is_list(self):
        """Test that NMS returns a Python list, not a numpy array."""
        boxes = np.array([[10.0, 10.0, 20.0, 20.0]], dtype=np.float64)
        scores = np.array([0.9], dtype=np.float64)
        codeflash_output = nms(boxes, scores, 0.5); result = codeflash_output # 46.1μs -> 40.9μs (12.6% faster)

    def test_return_indices_are_valid(self):
        """Test that returned indices are valid box indices."""
        boxes = np.array(
            [
                [10.0, 10.0, 20.0, 20.0],
                [50.0, 50.0, 60.0, 60.0],
                [100.0, 100.0, 110.0, 110.0],
            ],
            dtype=np.float64,
        )
        scores = np.array([0.9, 0.8, 0.7], dtype=np.float64)
        codeflash_output = nms(boxes, scores, 0.5); result = codeflash_output # 101μs -> 94.4μs (8.05% faster)
        # All returned indices should be within valid range
        for idx in result:
            pass

class TestNMSEdgeCases:
    """Test edge cases and boundary conditions."""

    def test_empty_input(self):
        """Test NMS with no bounding boxes."""
        boxes = np.empty((0, 4), dtype=np.float64)
        scores = np.empty(0, dtype=np.float64)
        codeflash_output = nms(boxes, scores, 0.5); result = codeflash_output # 17.3μs -> 17.5μs (1.26% slower)

    def test_zero_threshold(self):
        """Test NMS with threshold of 0 (very strict suppression)."""
        boxes = np.array(
            [[10.0, 10.0, 20.0, 20.0], [10.1, 10.1, 20.1, 20.1]], dtype=np.float64
        )
        scores = np.array([0.9, 0.8], dtype=np.float64)
        codeflash_output = nms(boxes, scores, 0.0); result = codeflash_output # 53.9μs -> 49.0μs (10.0% faster)

    def test_one_threshold(self):
        """Test NMS with threshold of 1.0 (no suppression)."""
        boxes = np.array(
            [[10.0, 10.0, 20.0, 20.0], [15.0, 15.0, 25.0, 25.0]], dtype=np.float64
        )
        scores = np.array([0.9, 0.8], dtype=np.float64)
        codeflash_output = nms(boxes, scores, 1.0); result = codeflash_output # 83.7μs -> 66.3μs (26.2% faster)

    def test_very_small_boxes(self):
        """Test NMS with very small bounding boxes."""
        boxes = np.array(
            [
                [0.0, 0.0, 0.1, 0.1],
                [0.05, 0.05, 0.15, 0.15],
                [1.0, 1.0, 1.1, 1.1],
            ],
            dtype=np.float64,
        )
        scores = np.array([0.9, 0.8, 0.7], dtype=np.float64)
        codeflash_output = nms(boxes, scores, 0.3); result = codeflash_output # 81.7μs -> 66.5μs (22.8% faster)

    def test_very_large_boxes(self):
        """Test NMS with very large bounding boxes."""
        boxes = np.array(
            [
                [0.0, 0.0, 10000.0, 10000.0],
                [5000.0, 5000.0, 15000.0, 15000.0],
                [20000.0, 20000.0, 30000.0, 30000.0],
            ],
            dtype=np.float64,
        )
        scores = np.array([0.9, 0.8, 0.7], dtype=np.float64)
        codeflash_output = nms(boxes, scores, 0.3); result = codeflash_output # 98.9μs -> 84.2μs (17.5% faster)

    def test_boxes_with_zero_area(self):
        """Test NMS with degenerate boxes (zero area)."""
        boxes = np.array(
            [
                [10.0, 10.0, 10.0, 10.0],  # zero area point
                [20.0, 20.0, 30.0, 30.0],  # normal box
            ],
            dtype=np.float64,
        )
        scores = np.array([0.9, 0.8], dtype=np.float64)
        codeflash_output = nms(boxes, scores, 0.5); result = codeflash_output # 83.2μs -> 65.6μs (26.9% faster)

    def test_negative_coordinates(self):
        """Test NMS with negative bounding box coordinates."""
        boxes = np.array(
            [
                [-10.0, -10.0, -5.0, -5.0],
                [-7.0, -7.0, -2.0, -2.0],
                [10.0, 10.0, 20.0, 20.0],
            ],
            dtype=np.float64,
        )
        scores = np.array([0.9, 0.8, 0.7], dtype=np.float64)
        codeflash_output = nms(boxes, scores, 0.3); result = codeflash_output # 101μs -> 93.7μs (8.55% faster)

    def test_perfect_overlap_different_scores(self):
        """Test NMS with perfectly overlapping boxes at different scores."""
        boxes = np.array(
            [
                [0.0, 0.0, 10.0, 10.0],
                [0.0, 0.0, 10.0, 10.0],
                [0.0, 0.0, 10.0, 10.0],
            ],
            dtype=np.float64,
        )
        scores = np.array([0.9, 0.7, 0.5], dtype=np.float64)
        codeflash_output = nms(boxes, scores, 0.5); result = codeflash_output # 53.8μs -> 47.4μs (13.6% faster)

    def test_single_coordinate_overlap(self):
        """Test NMS with boxes touching at a single point or edge."""
        boxes = np.array(
            [
                [0.0, 0.0, 10.0, 10.0],
                [10.0, 0.0, 20.0, 10.0],  # touching along edge
                [0.0, 10.0, 10.0, 20.0],  # touching along edge
            ],
            dtype=np.float64,
        )
        scores = np.array([0.9, 0.8, 0.7], dtype=np.float64)
        codeflash_output = nms(boxes, scores, 0.1); result = codeflash_output # 97.7μs -> 93.0μs (5.12% faster)

    def test_float32_input(self):
        """Test NMS with float32 input arrays."""
        boxes = np.array(
            [[10.0, 10.0, 20.0, 20.0], [50.0, 50.0, 60.0, 60.0]], dtype=np.float32
        )
        scores = np.array([0.9, 0.8], dtype=np.float32)
        codeflash_output = nms(boxes, scores, 0.5); result = codeflash_output # 81.5μs -> 73.2μs (11.3% faster)

    def test_scores_all_same(self):
        """Test NMS when all boxes have identical scores."""
        boxes = np.array(
            [
                [10.0, 10.0, 20.0, 20.0],
                [15.0, 15.0, 25.0, 25.0],
                [50.0, 50.0, 60.0, 60.0],
            ],
            dtype=np.float64,
        )
        scores = np.array([0.8, 0.8, 0.8], dtype=np.float64)
        codeflash_output = nms(boxes, scores, 0.3); result = codeflash_output # 101μs -> 92.3μs (10.4% faster)

    def test_inverted_box_coordinates(self):
        """Test NMS with inverted box coordinates (x2 < x1 or y2 < y1)."""
        boxes = np.array(
            [
                [20.0, 20.0, 10.0, 10.0],  # inverted
                [50.0, 50.0, 60.0, 60.0],  # normal
            ],
            dtype=np.float64,
        )
        scores = np.array([0.9, 0.8], dtype=np.float64)
        codeflash_output = nms(boxes, scores, 0.5); result = codeflash_output # 82.7μs -> 66.4μs (24.6% faster)

    def test_extreme_threshold_values(self):
        """Test NMS with extreme threshold values."""
        boxes = np.array(
            [[10.0, 10.0, 20.0, 20.0], [15.0, 15.0, 25.0, 25.0]], dtype=np.float64
        )
        scores = np.array([0.9, 0.8], dtype=np.float64)
        
        # Very small threshold
        codeflash_output = nms(boxes, scores, 1e-10); result_small = codeflash_output # 53.8μs -> 47.9μs (12.3% faster)
        
        # Threshold equal to 1.0
        codeflash_output = nms(boxes, scores, 1.0); result_large = codeflash_output # 51.2μs -> 46.3μs (10.8% faster)

class TestNMSLargeScale:
    """Test NMS with large-scale inputs to assess performance and scalability."""

    def test_many_non_overlapping_boxes(self):
        """Test NMS with many non-overlapping boxes."""
        # Create a grid of non-overlapping boxes
        n_boxes = 100
        boxes = []
        scores = []
        for i in range(n_boxes):
            row = i // 10
            col = i % 10
            x1, y1 = col * 100, row * 100
            x2, y2 = x1 + 50, y1 + 50
            boxes.append([x1, y1, x2, y2])
            scores.append(0.5 + 0.4 * (i / n_boxes))
        
        boxes = np.array(boxes, dtype=np.float64)
        scores = np.array(scores, dtype=np.float64)
        codeflash_output = nms(boxes, scores, 0.5); result = codeflash_output # 2.05ms -> 1.82ms (12.6% faster)

    def test_many_identical_boxes_different_scores(self):
        """Test NMS with many identical boxes having different scores."""
        n_boxes = 100
        boxes = np.tile(np.array([[10.0, 10.0, 20.0, 20.0]]), (n_boxes, 1))
        scores = np.linspace(0.1, 0.9, n_boxes)
        
        codeflash_output = nms(boxes, scores, 0.5); result = codeflash_output # 51.4μs -> 46.1μs (11.5% faster)

    def test_many_overlapping_boxes(self):
        """Test NMS with many highly overlapping boxes."""
        n_boxes = 50
        boxes = []
        scores = []
        
        # Create boxes that heavily overlap
        for i in range(n_boxes):
            offset = i * 2.0  # Small offset to ensure some variation
            x1 = 10.0 + offset
            y1 = 10.0 + offset
            x2 = 30.0 + offset
            y2 = 30.0 + offset
            boxes.append([x1, y1, x2, y2])
            scores.append(0.5 + 0.4 * (i / n_boxes))
        
        boxes = np.array(boxes, dtype=np.float64)
        scores = np.array(scores, dtype=np.float64)
        codeflash_output = nms(boxes, scores, 0.3); result = codeflash_output # 302μs -> 272μs (11.0% faster)

    def test_mixed_random_boxes(self):
        """Test NMS with random boxes (mixed overlapping and non-overlapping)."""
        np.random.seed(42)
        n_boxes = 200
        
        # Generate random boxes
        x1 = np.random.uniform(0, 800, n_boxes)
        y1 = np.random.uniform(0, 800, n_boxes)
        x2 = x1 + np.random.uniform(20, 100, n_boxes)
        y2 = y1 + np.random.uniform(20, 100, n_boxes)
        
        boxes = np.column_stack([x1, y1, x2, y2]).astype(np.float64)
        scores = np.random.uniform(0, 1, n_boxes).astype(np.float64)
        
        codeflash_output = nms(boxes, scores, 0.5); result = codeflash_output # 4.28ms -> 3.78ms (13.0% faster)

    def test_performance_moderate_scale(self):
        """Test NMS performance with a moderate number of boxes."""
        np.random.seed(42)
        n_boxes = 500
        
        x1 = np.random.uniform(0, 5000, n_boxes)
        y1 = np.random.uniform(0, 5000, n_boxes)
        x2 = x1 + np.random.uniform(50, 200, n_boxes)
        y2 = y1 + np.random.uniform(50, 200, n_boxes)
        
        boxes = np.column_stack([x1, y1, x2, y2]).astype(np.float64)
        scores = np.random.uniform(0, 1, n_boxes).astype(np.float64)
        
        codeflash_output = nms(boxes, scores, 0.5); result = codeflash_output # 13.5ms -> 12.1ms (11.9% faster)

    def test_clustered_boxes_large_scale(self):
        """Test NMS with large-scale clustered boxes (multiple clusters)."""
        np.random.seed(42)
        boxes = []
        scores = []
        
        # Create 5 clusters of boxes
        for cluster in range(5):
            cluster_x = cluster * 1000
            cluster_y = cluster * 1000
            
            for i in range(40):
                x1 = cluster_x + np.random.normal(100, 50)
                y1 = cluster_y + np.random.normal(100, 50)
                x2 = x1 + np.random.uniform(50, 150)
                y2 = y1 + np.random.uniform(50, 150)
                boxes.append([x1, y1, x2, y2])
                scores.append(np.random.uniform(0, 1))
        
        boxes = np.array(boxes, dtype=np.float64)
        scores = np.array(scores, dtype=np.float64)
        
        codeflash_output = nms(boxes, scores, 0.3); result = codeflash_output # 1.49ms -> 1.33ms (11.8% faster)

    def test_extreme_coordinate_range(self):
        """Test NMS with boxes spanning extreme coordinate ranges."""
        boxes = np.array(
            [
                [1e-6, 1e-6, 1e-5, 1e-5],  # Very small coordinates
                [1e6, 1e6, 1e6 + 100, 1e6 + 100],  # Very large coordinates
                [0.5, 0.5, 1.5, 1.5],  # Medium coordinates
                [1e3, 1e3, 1e4, 1e4],  # Different scale
            ],
            dtype=np.float64,
        )
        scores = np.array([0.9, 0.8, 0.7, 0.6], dtype=np.float64)
        
        codeflash_output = nms(boxes, scores, 0.5); result = codeflash_output # 122μs -> 112μs (8.82% faster)

    def test_result_is_sorted_by_input_order(self):
        """Test that NMS result maintains some consistency in box ordering."""
        np.random.seed(42)
        n_boxes = 100
        
        x1 = np.random.uniform(0, 1000, n_boxes)
        y1 = np.random.uniform(0, 1000, n_boxes)
        x2 = x1 + np.random.uniform(50, 150, n_boxes)
        y2 = y1 + np.random.uniform(50, 150, n_boxes)
        
        boxes = np.column_stack([x1, y1, x2, y2]).astype(np.float64)
        scores = np.random.uniform(0, 1, n_boxes).astype(np.float64)
        
        codeflash_output = nms(boxes, scores, 0.5); result = codeflash_output # 1.99ms -> 1.75ms (13.3% faster)

class TestNMSRobustness:
    """Test robustness and consistency of the NMS function."""

    def test_deterministic_results(self):
        """Test that NMS produces deterministic results for the same input."""
        boxes = np.array(
            [
                [10.0, 10.0, 20.0, 20.0],
                [15.0, 15.0, 25.0, 25.0],
                [50.0, 50.0, 60.0, 60.0],
            ],
            dtype=np.float64,
        )
        scores = np.array([0.9, 0.8, 0.7], dtype=np.float64)
        
        codeflash_output = nms(boxes, scores, 0.3); result1 = codeflash_output # 103μs -> 94.3μs (9.21% faster)
        codeflash_output = nms(boxes, scores, 0.3); result2 = codeflash_output # 68.4μs -> 60.3μs (13.4% faster)

    def test_no_duplicate_indices(self):
        """Test that NMS never returns duplicate indices."""
        np.random.seed(42)
        n_boxes = 100
        
        x1 = np.random.uniform(0, 500, n_boxes)
        y1 = np.random.uniform(0, 500, n_boxes)
        x2 = x1 + np.random.uniform(20, 100, n_boxes)
        y2 = y1 + np.random.uniform(20, 100, n_boxes)
        
        boxes = np.column_stack([x1, y1, x2, y2]).astype(np.float64)
        scores = np.random.uniform(0, 1, n_boxes).astype(np.float64)
        
        codeflash_output = nms(boxes, scores, 0.5); result = codeflash_output # 1.98ms -> 1.74ms (14.0% faster)

    def test_indices_within_bounds(self):
        """Test that all returned indices are within valid bounds."""
        np.random.seed(42)
        
        for n_boxes in [10, 50, 100]:
            x1 = np.random.uniform(0, 500, n_boxes)
            y1 = np.random.uniform(0, 500, n_boxes)
            x2 = x1 + np.random.uniform(20, 100, n_boxes)
            y2 = y1 + np.random.uniform(20, 100, n_boxes)
            
            boxes = np.column_stack([x1, y1, x2, y2]).astype(np.float64)
            scores = np.random.uniform(0, 1, n_boxes).astype(np.float64)
            
            codeflash_output = nms(boxes, scores, 0.5); result = codeflash_output # 3.09ms -> 2.73ms (13.2% faster)
            
            for idx in result:
                pass

    def test_highest_score_always_kept(self):
        """Test that the box with highest score is always kept (when no overlap threshold)."""
        boxes = np.array(
            [
                [10.0, 10.0, 20.0, 20.0],
                [50.0, 50.0, 60.0, 60.0],
                [100.0, 100.0, 110.0, 110.0],
            ],
            dtype=np.float64,
        )
        scores = np.array([0.5, 0.3, 0.9], dtype=np.float64)
        codeflash_output = nms(boxes, scores, 0.5); result = codeflash_output # 103μs -> 93.8μs (10.4% faster)

    def test_contiguous_array_handling(self):
        """Test that NMS handles both contiguous and non-contiguous arrays."""
        boxes_c = np.array(
            [[10.0, 10.0, 20.0, 20.0], [50.0, 50.0, 60.0, 60.0]], dtype=np.float64
        )
        scores_c = np.array([0.9, 0.8], dtype=np.float64)
        
        codeflash_output = nms(boxes_c, scores_c, 0.5); result_c = codeflash_output # 73.5μs -> 65.9μs (11.5% faster)
        
        # Create non-contiguous view by transposing and selecting columns
        boxes_nc = boxes_c.T[::1].T
        scores_nc = scores_c[::1]
        
        codeflash_output = nms(boxes_nc, scores_nc, 0.5); result_nc = codeflash_output # 49.5μs -> 44.1μs (12.1% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-nms-mkowazdu and push.