⚡️ Speed up function find_supported_resolutions by 16%

[codeflash-ai]

📄 16% (0.16x) speedup for find_supported_resolutions in src/transformers/models/llama4/image_processing_llama4_fast.py

⏱️ Runtime : 14.9 milliseconds → 12.8 milliseconds (best of 141 runs)

📝 Explanation and details

The optimized code achieves a 15% speedup through several key algorithmic and implementation improvements:

1. Precomputed square root in get_factors:
The most significant optimization is extracting int(dividend**0.5) outside the loop and storing it in the limit variable. This eliminates thousands of redundant square root calculations - the line profiler shows this loop runs ~74,443 times across all test cases. Computing the square root once instead of on every iteration reduces computational overhead.

2. List comprehension for resolution generation:
The nested loop structure for building possible_resolutions was replaced with a single list comprehension:

possible_resolutions = [
    (height * patch_size_val, width * patch_size_val)
    for value in asp_dict.values()
    for height, width in value
]

This eliminates the overhead of repeated append() calls and leverages Python's optimized list comprehension implementation.

3. Variable renaming for clarity:
The patch_size variable was renamed to patch_size_val to avoid name shadowing with the input parameter, which can cause subtle performance impacts due to namespace lookups.

Performance characteristics:

• Large-scale scenarios benefit most: Tests with high max_num_chunks (e.g., 999+ chunks) show 16-18% improvements because they maximize the impact of the precomputed square root optimization
• Small-scale scenarios see moderate gains: Basic cases with few chunks show 6-15% improvements primarily from the list comprehension optimization
• Edge cases maintain correctness: All edge cases (zero/negative chunks, non-square patches) preserve identical behavior while gaining performance benefits

The optimizations are particularly effective for image processing workloads where find_supported_resolutions is called with large chunk counts, as the factor computation dominates the runtime.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from collections import namedtuple

# imports
import pytest
import torch
from transformers.models.llama4.image_processing_llama4_fast import \
    find_supported_resolutions

# Define a minimal SizeDict for testing
SizeDict = namedtuple("SizeDict", ["height", "width"])
from transformers.models.llama4.image_processing_llama4_fast import \
    find_supported_resolutions

# ------------------- UNIT TESTS -------------------

# Basic Test Cases

def test_basic_1_chunk_1_patch():
    # Only one chunk, patch size 1x1
    patch_size = SizeDict(height=1, width=1)
    codeflash_output = find_supported_resolutions(1, patch_size); result = codeflash_output # 7.62μs -> 7.16μs (6.41% faster)

def test_basic_2_chunks_patch_2():
    # Two chunks, patch size 2x2
    patch_size = SizeDict(height=2, width=2)
    codeflash_output = find_supported_resolutions(2, patch_size); result = codeflash_output # 10.3μs -> 9.49μs (8.05% faster)
    # Factors of 2: 1, 2
    # Possible: (1,2), (2,1), (1,1)
    expected = torch.tensor([[2, 4], [4, 2], [2, 2]])
    # All expected resolutions must be present
    for res in expected:
        pass

def test_basic_5_chunks_patch_224():
    # Example from docstring
    patch_size = SizeDict(height=224, width=224)
    codeflash_output = find_supported_resolutions(5, patch_size); result = codeflash_output # 16.2μs -> 15.1μs (7.42% faster)
    # Factors of 5: 1, 5
    # For each chunk_size from 5 down to 1
    # Just check some expected resolutions
    expected = [
        (224, 1120),  # 1x5
        (1120, 224),  # 5x1
        (224, 896),   # 1x4
        (448, 448),   # 2x2
        (224, 224),   # 1x1
    ]
    for res in expected:
        pass

def test_basic_patch_size_10_chunks_6():
    # patch_size=10, max_num_chunks=6
    patch_size = SizeDict(height=10, width=10)
    codeflash_output = find_supported_resolutions(6, patch_size); result = codeflash_output # 18.2μs -> 16.8μs (8.31% faster)
    # Factors for 6: 1,2,3,6
    expected = [
        (10, 60), (20, 30), (30, 20), (60, 10),  # chunk_size=6
        (10, 50), (50, 10), (10, 40), (20, 20), (40, 10),  # chunk_size=5,4
        (10, 30), (30, 10), (10, 20), (20, 10), (10, 10),  # chunk_size=3,2,1
    ]
    for res in expected:
        pass

# Edge Test Cases

def test_edge_patch_size_not_square():
    # Should raise ValueError if patch is not square
    patch_size = SizeDict(height=10, width=20)
    with pytest.raises(ValueError):
        find_supported_resolutions(4, patch_size) # 1.46μs -> 1.47μs (0.952% slower)

def test_edge_zero_chunks():
    # Zero chunks: should return empty tensor
    patch_size = SizeDict(height=10, width=10)
    codeflash_output = find_supported_resolutions(0, patch_size); result = codeflash_output # 3.39μs -> 3.77μs (9.98% slower)

def test_edge_negative_chunks():
    # Negative chunks: should return empty tensor
    patch_size = SizeDict(height=10, width=10)
    codeflash_output = find_supported_resolutions(-5, patch_size); result = codeflash_output # 3.19μs -> 3.57μs (10.6% slower)

def test_edge_patch_size_zero():
    # Patch size zero: all resolutions are zero
    patch_size = SizeDict(height=0, width=0)
    codeflash_output = find_supported_resolutions(3, patch_size); result = codeflash_output # 13.1μs -> 11.9μs (10.3% faster)

def test_edge_patch_size_negative():
    # Negative patch size: resolutions are negative
    patch_size = SizeDict(height=-8, width=-8)
    codeflash_output = find_supported_resolutions(2, patch_size); result = codeflash_output # 10.6μs -> 9.38μs (12.7% faster)
    # All resolutions should be negative multiples of -8
    for res in result:
        pass

def test_edge_large_patch_size_small_chunks():
    # Large patch, small chunk count
    patch_size = SizeDict(height=999, width=999)
    codeflash_output = find_supported_resolutions(1, patch_size); result = codeflash_output # 7.54μs -> 7.14μs (5.61% faster)


def test_large_scale_max_chunks_1000_patch_1():
    # max_num_chunks=1000, patch_size=1
    patch_size = SizeDict(height=1, width=1)
    codeflash_output = find_supported_resolutions(1000, patch_size); result = codeflash_output # 4.05ms -> 3.53ms (14.6% faster)
    # All resolutions should be <= 1000
    for res in result:
        pass

def test_large_scale_patch_size_100_chunks_100():
    # max_num_chunks=100, patch_size=100
    patch_size = SizeDict(height=100, width=100)
    codeflash_output = find_supported_resolutions(100, patch_size); result = codeflash_output # 291μs -> 254μs (14.6% faster)
    # All resolutions should be multiples of 100
    for res in result:
        pass
    # Should not exceed (100*100, 100*100)
    for res in result:
        pass

def test_large_scale_patch_size_500_chunks_10():
    # max_num_chunks=10, patch_size=500
    patch_size = SizeDict(height=500, width=500)
    codeflash_output = find_supported_resolutions(10, patch_size); result = codeflash_output # 28.5μs -> 24.8μs (14.8% faster)
    # All resolutions should be multiples of 500
    for res in result:
        pass
    # Should not exceed (5000, 5000)
    for res in result:
        pass

def test_large_scale_patch_size_99_chunks_999():
    # max_num_chunks=999, patch_size=99
    patch_size = SizeDict(height=99, width=99)
    codeflash_output = find_supported_resolutions(999, patch_size); result = codeflash_output # 4.19ms -> 3.61ms (16.1% faster)
    # All resolutions should be multiples of 99
    for res in result:
        pass
    # Should not exceed (99*999, 99*999)
    for res in result:
        pass

# Additional edge: very large patch, small chunks, ensure memory safety
def test_large_patch_small_chunks_memory():
    # patch_size=10000, max_num_chunks=2
    patch_size = SizeDict(height=10000, width=10000)
    codeflash_output = find_supported_resolutions(2, patch_size); result = codeflash_output # 12.1μs -> 11.3μs (6.86% faster)
    # Resolutions are correct
    expected = torch.tensor([[10000, 20000], [20000, 10000], [10000, 10000]])
    for res in expected:
        pass
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from collections import defaultdict
from functools import lru_cache

# imports
import pytest  # used for our unit tests
# function to test
import torch
from transformers.models.llama4.image_processing_llama4_fast import \
    find_supported_resolutions


class SizeDict:
    """Minimal implementation for testing."""
    def __init__(self, height, width):
        self.height = height
        self.width = width
from transformers.models.llama4.image_processing_llama4_fast import \
    find_supported_resolutions

# unit tests

# --- BASIC TEST CASES ---

def test_basic_1_chunk_patch_224():
    # Only one chunk, patch size 224
    patch = SizeDict(224, 224)
    codeflash_output = find_supported_resolutions(1, patch); result = codeflash_output # 8.25μs -> 7.94μs (3.87% faster)

def test_basic_2_chunks_patch_224():
    # Two chunks, patch size 224
    patch = SizeDict(224, 224)
    codeflash_output = find_supported_resolutions(2, patch); result = codeflash_output # 10.8μs -> 10.0μs (7.64% faster)
    # Possible splits: (1,2), (2,1)
    expected = [(224, 448), (448, 224)]
    for res in expected:
        pass

def test_basic_4_chunks_patch_224():
    # Four chunks, patch size 224
    patch = SizeDict(224, 224)
    codeflash_output = find_supported_resolutions(4, patch); result = codeflash_output # 15.0μs -> 13.0μs (14.8% faster)
    # Possible splits: (1,4), (2,2), (4,1)
    expected = [(224, 896), (448, 448), (896, 224)]
    for res in expected:
        pass

def test_basic_patch_size_32():
    # Patch size 32, 3 chunks
    patch = SizeDict(32, 32)
    codeflash_output = find_supported_resolutions(3, patch); result = codeflash_output # 11.9μs -> 10.7μs (11.2% faster)
    # Possible splits: (1,3), (3,1)
    expected = [(32, 96), (96, 32)]
    for res in expected:
        pass

def test_basic_patch_size_1():
    # Patch size 1, 2 chunks
    patch = SizeDict(1, 1)
    codeflash_output = find_supported_resolutions(2, patch); result = codeflash_output # 9.94μs -> 8.98μs (10.7% faster)
    expected = [(1, 2), (2, 1)]
    for res in expected:
        pass

# --- EDGE TEST CASES ---

def test_edge_patch_size_not_square():
    # Patch size not square should raise ValueError
    patch = SizeDict(224, 112)
    with pytest.raises(ValueError):
        find_supported_resolutions(2, patch) # 1.40μs -> 1.39μs (0.144% faster)

def test_edge_max_num_chunks_zero():
    # Zero chunks: should return empty list
    patch = SizeDict(224, 224)
    codeflash_output = find_supported_resolutions(0, patch); result = codeflash_output # 3.31μs -> 3.61μs (8.37% slower)

def test_edge_max_num_chunks_negative():
    # Negative chunks: should return empty list
    patch = SizeDict(224, 224)
    codeflash_output = find_supported_resolutions(-1, patch); result = codeflash_output # 3.06μs -> 3.44μs (11.0% slower)

def test_edge_patch_size_zero():
    # Patch size zero: resolutions will be zero
    patch = SizeDict(0, 0)
    codeflash_output = find_supported_resolutions(2, patch); result = codeflash_output # 11.0μs -> 10.2μs (8.68% faster)

def test_edge_patch_size_negative():
    # Negative patch size: resolutions will be negative
    patch = SizeDict(-16, -16)
    codeflash_output = find_supported_resolutions(3, patch); result = codeflash_output # 12.2μs -> 11.0μs (10.8% faster)
    for h, w in result:
        pass

def test_edge_max_num_chunks_one():
    # Only one chunk, patch size 10
    patch = SizeDict(10, 10)
    codeflash_output = find_supported_resolutions(1, patch); result = codeflash_output # 7.46μs -> 6.99μs (6.74% faster)


def test_edge_large_patch_size():
    # Large patch size, small chunk
    patch = SizeDict(999, 999)
    codeflash_output = find_supported_resolutions(2, patch); result = codeflash_output # 11.9μs -> 10.8μs (9.41% faster)
    expected = [(999, 1998), (1998, 999)]
    for res in expected:
        pass

# --- LARGE SCALE TEST CASES ---



def test_large_scale_max_num_chunks_999_patch_1():
    # Max possible chunks (999), patch size 1
    patch = SizeDict(1, 1)
    codeflash_output = find_supported_resolutions(999, patch); result = codeflash_output # 4.10ms -> 3.52ms (16.6% faster)

def test_large_scale_unique_resolutions():
    # Check that all resolutions are unique
    patch = SizeDict(5, 5)
    codeflash_output = find_supported_resolutions(50, patch); result = codeflash_output # 136μs -> 115μs (18.1% faster)

def test_large_scale_performance():
    # Performance: Should not take long for 500 chunks
    import time
    patch = SizeDict(2, 2)
    start = time.time()
    codeflash_output = find_supported_resolutions(500, patch); result = codeflash_output # 1.79ms -> 1.51ms (18.7% faster)
    elapsed = time.time() - start

# --- DETERMINISM TEST CASE ---

def test_determinism():
    # Calling twice should yield same result
    patch = SizeDict(32, 32)
    codeflash_output = find_supported_resolutions(7, patch); r1 = codeflash_output # 21.0μs -> 18.7μs (12.4% faster)
    codeflash_output = find_supported_resolutions(7, patch); r2 = codeflash_output # 264ns -> 271ns (2.58% slower)

# --- FUNCTIONALITY TEST CASES ---

def test_resolution_values_are_multiples_of_patch_size():
    # All returned resolutions should be multiples of patch size
    patch = SizeDict(20, 20)
    codeflash_output = find_supported_resolutions(6, patch); result = codeflash_output # 17.8μs -> 15.6μs (14.6% faster)
    for h, w in result:
        pass

def test_all_possible_aspect_ratios_included():
    # For chunks=6, patch=10, aspect ratios should include 1:6, 2:3, 3:2, 6:1, etc.
    patch = SizeDict(10, 10)
    codeflash_output = find_supported_resolutions(6, patch); result = codeflash_output # 17.2μs -> 15.4μs (11.9% faster)
    expected = {(10, 60), (20, 30), (30, 20), (60, 10)}
    for res in expected:
        pass
# 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-find_supported_resolutions-mhjqxcgk and push.