⚡️ Speed up function compute_intermediate_size by 13%

[codeflash-ai]

📄 13% (0.13x) speedup for compute_intermediate_size in src/transformers/models/olmo2/convert_olmo2_weights_to_hf.py

⏱️ Runtime : 646 microseconds → 570 microseconds (best of 171 runs)

📝 Explanation and details

The optimization replaces int(8 * n / 3) with (8 * n // 3) in the computation, achieving a 13% speedup by eliminating unnecessary floating-point operations.

Key optimization:

• Original: int(ffn_dim_multiplier * int(8 * n / 3)) performs floating-point division (/) then converts to int
• Optimized: int(ffn_dim_multiplier * (8 * n // 3)) uses integer floor division (//) directly

Why this is faster:

• Integer floor division (//) operates entirely in integer arithmetic, avoiding the overhead of converting to float and back to int
• The / operator in Python creates a float intermediate result that must be cast back with int(), adding unnecessary computation
• Since we only need the integer quotient, // is the more direct and efficient operation

Performance characteristics:

• The optimization shows consistent 10-30% improvements across most test cases involving integer inputs
• Particularly effective for basic computations with default parameters (25-34% faster)
• Some edge cases with float inputs show minor slowdowns due to type conversion overhead, but this represents the minority of real-world usage

Mathematical equivalence:
Both expressions produce identical results since int(8 * n / 3) and (8 * n // 3) yield the same integer quotient for all integer inputs, preserving all functional behavior while improving performance through more efficient arithmetic operations.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

# imports
import pytest  # used for our unit tests

from transformers.models.olmo2.convert_olmo2_weights_to_hf import compute_intermediate_size


# unit tests

# ---------------------------
# Basic Test Cases
# ---------------------------


def test_basic_default_params():
    # Basic scenario with default parameters
    # n = 768, ffn_dim_multiplier=1, multiple_of=256
    # Expected: int(8*768/3) = 2048, so output = 2048
    codeflash_output = compute_intermediate_size(768)  # 1.12μs -> 833ns (34.2% faster)


def test_basic_non_default_multiplier():
    # n = 768, ffn_dim_multiplier=1.5, multiple_of=256
    # int(8*768/3) = 2048, *1.5 = 3072, output = 3072
    codeflash_output = compute_intermediate_size(768, ffn_dim_multiplier=1.5)  # 1.36μs -> 1.22μs (12.0% faster)


def test_basic_non_default_multiple_of():
    # n = 768, ffn_dim_multiplier=1, multiple_of=512
    # int(8*768/3) = 2048, output = 2048 rounded up to nearest 512 = 2048
    codeflash_output = compute_intermediate_size(768, multiple_of=512)  # 1.30μs -> 1.05μs (24.2% faster)


def test_basic_rounding_up():
    # n = 700, ffn_dim_multiplier=1, multiple_of=256
    # int(8*700/3) = 1866, rounded up to 2048
    codeflash_output = compute_intermediate_size(700)  # 1.08μs -> 833ns (29.3% faster)


def test_basic_rounding_down():
    # n = 256, ffn_dim_multiplier=1, multiple_of=256
    # int(8*256/3) = 682, rounded up to 768
    codeflash_output = compute_intermediate_size(256)  # 1.02μs -> 811ns (25.6% faster)


# ---------------------------
# Edge Test Cases
# ---------------------------


def test_edge_n_zero():
    # n=0 should always return 0
    codeflash_output = compute_intermediate_size(0)  # 1.01μs -> 908ns (10.8% faster)


def test_edge_n_one():
    # n=1, int(8*1/3) = 2, rounded up to 256
    codeflash_output = compute_intermediate_size(1)  # 1.01μs -> 829ns (21.6% faster)


def test_edge_multiplier_zero():
    # ffn_dim_multiplier=0, should always return 0
    codeflash_output = compute_intermediate_size(768, ffn_dim_multiplier=0)  # 1.38μs -> 1.10μs (25.0% faster)


def test_edge_multiplier_fractional():
    # n=768, ffn_dim_multiplier=0.5, int(8*768/3)=2048, *0.5=1024, rounded up to 1024
    codeflash_output = compute_intermediate_size(768, ffn_dim_multiplier=0.5)  # 1.42μs -> 1.29μs (10.4% faster)


def test_edge_multiple_of_one():
    # n=768, multiple_of=1, int(8*768/3)=2048, output=2048
    codeflash_output = compute_intermediate_size(768, multiple_of=1)  # 1.24μs -> 1.00μs (23.9% faster)


def test_edge_multiple_of_large():
    # n=768, multiple_of=4096, int(8*768/3)=2048, rounded up to 4096
    codeflash_output = compute_intermediate_size(768, multiple_of=4096)  # 1.27μs -> 1.02μs (24.7% faster)


def test_edge_multiple_of_not_divisor():
    # n=100, multiple_of=7, int(8*100/3)=266, rounded up to 273 (7*39)
    codeflash_output = compute_intermediate_size(100, multiple_of=7)  # 1.28μs -> 967ns (32.3% faster)


def test_edge_negative_n():
    # Negative n should still compute, but result is 0 (since int(8*-1/3)=-2, max(0,...) is not used, but rounding up to 0)
    codeflash_output = compute_intermediate_size(-1)  # 1.07μs -> 856ns (25.5% faster)


def test_edge_negative_multiplier():
    # Negative multiplier should produce 0 for positive n (since int(8*768/3)=2048, *-1=-2048, rounded up to 0)
    codeflash_output = compute_intermediate_size(768, ffn_dim_multiplier=-1)  # 1.49μs -> 1.17μs (27.0% faster)


def test_edge_negative_multiple_of():
    # Negative multiple_of, e.g. -256, should produce 0 or negative multiples
    # n=768, int(8*768/3)=2048, rounded up to next lowest negative multiple: -256 * ((2048 + -256 - 1)//-256)
    # ((2048 - 257)//-256) = (1791//-256) = -7, -256*-7 = 1792
    # But negative multiple_of is a strange case, so let's check the actual calculation
    codeflash_output = compute_intermediate_size(768, multiple_of=-256)
    result = codeflash_output  # 1.35μs -> 1.02μs (32.3% faster)


def test_edge_large_fractional_multiplier():
    # n=768, ffn_dim_multiplier=2.7, int(8*768/3)=2048, *2.7=5529.6, int=5529, rounded up to 5632
    codeflash_output = compute_intermediate_size(768, ffn_dim_multiplier=2.7)  # 1.43μs -> 1.28μs (11.4% faster)


def test_edge_float_n():
    # n=768.0 (float), should behave same as int
    codeflash_output = compute_intermediate_size(768.0)  # 1.29μs -> 1.64μs (21.3% slower)


def test_edge_float_multiple_of():
    # n=768, multiple_of=256.0 (float), should behave same as int
    codeflash_output = compute_intermediate_size(768, multiple_of=256.0)  # 1.82μs -> 1.53μs (18.5% faster)


def test_edge_float_multiplier():
    # n=768, ffn_dim_multiplier=1.0 (float), should behave same as int
    codeflash_output = compute_intermediate_size(768, ffn_dim_multiplier=1.0)  # 1.40μs -> 1.20μs (16.9% faster)


def test_edge_small_multiple_of():
    # n=7, multiple_of=2, int(8*7/3)=18, rounded up to 18
    codeflash_output = compute_intermediate_size(7, multiple_of=2)  # 1.16μs -> 931ns (24.8% faster)


def test_edge_non_integer_result():
    # n=5, int(8*5/3)=13, rounded up to 256
    codeflash_output = compute_intermediate_size(5)  # 983ns -> 828ns (18.7% faster)


# ---------------------------
# Large Scale Test Cases
# ---------------------------


def test_large_n():
    # Large n, e.g. n=999, default params
    # int(8*999/3)=2664, rounded up to 2816
    codeflash_output = compute_intermediate_size(999)  # 1.12μs -> 858ns (30.2% faster)


def test_large_multiplier():
    # Large ffn_dim_multiplier
    # n=512, ffn_dim_multiplier=10, int(8*512/3)=1365, *10=13650, rounded up to 13824
    codeflash_output = compute_intermediate_size(512, ffn_dim_multiplier=10)  # 1.23μs -> 1.08μs (14.1% faster)


def test_large_multiple_of():
    # Large multiple_of
    # n=512, multiple_of=999, int(8*512/3)=1365, rounded up to 1998
    codeflash_output = compute_intermediate_size(512, multiple_of=999)  # 1.26μs -> 1.01μs (24.6% faster)


def test_large_all_params():
    # n=999, ffn_dim_multiplier=5.5, multiple_of=999
    # int(8*999/3)=2664, *5.5=14652, rounded up to 14985
    codeflash_output = compute_intermediate_size(
        999, ffn_dim_multiplier=5.5, multiple_of=999
    )  # 1.47μs -> 1.27μs (15.8% faster)


def test_large_scale_stress():
    # Test with n from 1 to 1000, ensure no exceptions and outputs are multiples of 256
    for n in range(1, 1001):
        codeflash_output = compute_intermediate_size(n)
        result = codeflash_output  # 254μs -> 222μs (14.3% faster)


def test_large_scale_multiplier_stress():
    # Test with multipliers from 0.1 to 10.0 in steps of 0.1 for n=512
    for m in [round(x * 0.1, 2) for x in range(1, 101)]:
        codeflash_output = compute_intermediate_size(512, ffn_dim_multiplier=m)
        result = codeflash_output  # 31.9μs -> 28.6μs (11.5% faster)


def test_large_scale_multiple_of_stress():
    # Test with multiple_of from 1 to 1000 in steps of 111 for n=512
    for mo in range(1, 1001, 111):
        codeflash_output = compute_intermediate_size(512, multiple_of=mo)
        result = codeflash_output  # 4.03μs -> 3.52μs (14.7% faster)


# ---------------------------
# Determinism Test
# ---------------------------


def test_determinism():
    # Given same inputs, function should always return same output
    for n in [1, 100, 512, 999]:
        for m in [1, 2, 3.5]:
            for mo in [1, 256, 512]:
                codeflash_output = compute_intermediate_size(n, ffn_dim_multiplier=m, multiple_of=mo)
                result1 = codeflash_output
                codeflash_output = compute_intermediate_size(n, ffn_dim_multiplier=m, multiple_of=mo)
                result2 = codeflash_output


# ---------------------------
# Type Robustness Test
# ---------------------------


def test_type_robustness():
    # Should accept both int and float for all parameters
    codeflash_output = compute_intermediate_size(
        768.0, ffn_dim_multiplier=1.0, multiple_of=256.0
    )  # 1.95μs -> 1.74μs (12.3% faster)
    codeflash_output = compute_intermediate_size(
        768, ffn_dim_multiplier=1.0, multiple_of=256
    )  # 611ns -> 616ns (0.812% slower)
    codeflash_output = compute_intermediate_size(
        768.0, ffn_dim_multiplier=1, multiple_of=256.0
    )  # 519ns -> 489ns (6.13% faster)


# ---------------------------
# Exception/Invalid Input Test
# ---------------------------


def test_invalid_types():
    # Should raise TypeError for completely invalid types
    with pytest.raises(TypeError):
        compute_intermediate_size("768")
    with pytest.raises(TypeError):
        compute_intermediate_size(768, ffn_dim_multiplier="1")
    with pytest.raises(TypeError):
        compute_intermediate_size(768, multiple_of="256")


def test_invalid_none():
    # Should raise TypeError for None
    with pytest.raises(TypeError):
        compute_intermediate_size(None)  # 1.73μs -> 1.65μs (5.10% faster)
    with pytest.raises(TypeError):
        compute_intermediate_size(768, ffn_dim_multiplier=None)  # 1.56μs -> 1.32μs (18.1% faster)
    with pytest.raises(TypeError):
        compute_intermediate_size(768, multiple_of=None)  # 1.09μs -> 1.01μs (7.83% faster)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import pytest  # used for our unit tests

from transformers.models.olmo2.convert_olmo2_weights_to_hf import compute_intermediate_size


# unit tests

# -------------------- Basic Test Cases --------------------


def test_basic_default_params():
    # Test with n=768, default multiplier and multiple_of
    # int(8*768/3) = int(2048) = 2048
    # 2048*1 = 2048, round up to nearest multiple of 256 = 2048
    codeflash_output = compute_intermediate_size(768)  # 1.24μs -> 982ns (25.8% faster)


def test_basic_non_default_multiplier():
    # Test with n=768, multiplier=1.5, default multiple_of
    # int(8*768/3) = 2048, 2048*1.5=3072, round up to 3072
    codeflash_output = compute_intermediate_size(768, ffn_dim_multiplier=1.5)  # 1.39μs -> 1.24μs (12.8% faster)


def test_basic_non_default_multiple_of():
    # Test with n=768, default multiplier, multiple_of=512
    # int(8*768/3) = 2048, round up to nearest 512 = 2048
    codeflash_output = compute_intermediate_size(768, multiple_of=512)  # 1.17μs -> 1.01μs (16.2% faster)


def test_basic_both_non_default():
    # Test with n=600, multiplier=1.33, multiple_of=128
    # int(8*600/3) = 1600, 1600*1.33=2128, int(2128)=2128, round up to 2176
    codeflash_output = compute_intermediate_size(
        600, ffn_dim_multiplier=1.33, multiple_of=128
    )  # 1.35μs -> 1.19μs (14.1% faster)


def test_basic_small_n():
    # Test with small n=3, default multiplier and multiple_of
    # int(8*3/3) = int(8) = 8, round up to 256
    codeflash_output = compute_intermediate_size(3)  # 982ns -> 831ns (18.2% faster)


# -------------------- Edge Test Cases --------------------


def test_edge_n_zero():
    # Test with n=0, should always round up to 0
    codeflash_output = compute_intermediate_size(0)  # 961ns -> 856ns (12.3% faster)


def test_edge_n_one():
    # Test with n=1, int(8*1/3)=2, round up to 256
    codeflash_output = compute_intermediate_size(1)  # 960ns -> 795ns (20.8% faster)


def test_edge_n_negative():
    # Test with n=-10, should handle negative gracefully
    # int(8*-10/3) = int(-26.666...) = -26, round up to 0
    codeflash_output = compute_intermediate_size(-10)  # 1.21μs -> 914ns (32.7% faster)


def test_edge_multiplier_zero():
    # Test with multiplier=0, should always round up to 0
    codeflash_output = compute_intermediate_size(512, ffn_dim_multiplier=0)  # 1.26μs -> 1.03μs (22.0% faster)


def test_edge_multiplier_negative():
    # Test with negative multiplier, e.g. -1
    # int(8*512/3)=1365, 1365*-1=-1365, round up to 0
    codeflash_output = compute_intermediate_size(512, ffn_dim_multiplier=-1)  # 1.33μs -> 1.14μs (16.8% faster)


def test_edge_multiple_of_one():
    # Test with multiple_of=1, should not round up
    # int(8*100/3)=266, 266*1=266, round up to 266
    codeflash_output = compute_intermediate_size(100, multiple_of=1)  # 1.24μs -> 1.00μs (24.0% faster)


def test_edge_multiple_of_not_divisor():
    # Test with a multiple_of not a divisor of result
    # n=100, int(8*100/3)=266, 266*1.1=292.6, int=292, round up to 300 (next 10)
    codeflash_output = compute_intermediate_size(
        100, ffn_dim_multiplier=1.1, multiple_of=10
    )  # 1.39μs -> 1.27μs (9.26% faster)


def test_edge_multiple_of_large():
    # Test with large multiple_of
    # n=1000, int(8*1000/3)=2666, round up to 5000
    codeflash_output = compute_intermediate_size(1000, multiple_of=5000)  # 1.22μs -> 992ns (22.8% faster)


def test_edge_non_integer_multiplier():
    # Test with non-integer multiplier
    # n=100, int(8*100/3)=266, 266*1.25=332.5, int=332, round up to 384
    codeflash_output = compute_intermediate_size(
        100, ffn_dim_multiplier=1.25, multiple_of=64
    )  # 1.33μs -> 1.20μs (11.1% faster)


def test_edge_non_integer_n():
    # Test with non-integer n, e.g. n=123.7
    # int(8*123.7/3)=int(329.866...) = 329, 329*1=329, round up to 384
    codeflash_output = compute_intermediate_size(123.7, multiple_of=64)  # 1.51μs -> 1.82μs (17.1% slower)


def test_edge_large_multiplier():
    # Test with very large multiplier
    # n=10, int(8*10/3)=26, 26*100=2600, round up to 2560
    codeflash_output = compute_intermediate_size(
        10, ffn_dim_multiplier=100, multiple_of=256
    )  # 1.20μs -> 1.03μs (16.6% faster)


def test_edge_large_negative_n():
    # Test with large negative n
    codeflash_output = compute_intermediate_size(-10000)  # 1.22μs -> 917ns (32.7% faster)


# -------------------- Large Scale Test Cases --------------------


def test_large_n():
    # Test with large n
    # n=100000, int(8*100000/3)=266666, round up to 266752
    codeflash_output = compute_intermediate_size(100000, multiple_of=64)  # 1.26μs -> 1.03μs (22.1% faster)


def test_large_n_and_multiplier():
    # Test with large n and multiplier
    # n=999999, int(8*999999/3)=2666664, *1.5=3999996, round up to 4000000
    codeflash_output = compute_intermediate_size(
        999999, ffn_dim_multiplier=1.5, multiple_of=8
    )  # 1.43μs -> 1.23μs (16.0% faster)


def test_large_multiple_of():
    # Test with large multiple_of
    # n=1000, int(8*1000/3)=2666, round up to 10000
    codeflash_output = compute_intermediate_size(1000, multiple_of=10000)  # 1.24μs -> 980ns (26.8% faster)


def test_large_all_params():
    # n=999, multiplier=2.5, multiple_of=512
    # int(8*999/3)=2664, *2.5=6660, round up to 6912
    codeflash_output = compute_intermediate_size(
        999, ffn_dim_multiplier=2.5, multiple_of=512
    )  # 1.32μs -> 1.21μs (9.08% faster)


def test_large_batch_of_ns():
    # Test a batch of n values for monotonicity and rounding
    for n in range(1, 1000, 100):  # n = 1, 101, 201, ...
        codeflash_output = compute_intermediate_size(n)
        val = codeflash_output  # 3.62μs -> 3.14μs (15.5% faster)


def test_large_batch_monotonicity():
    # Test that increasing n does not decrease output
    prev = 0
    for n in range(1, 1000):
        codeflash_output = compute_intermediate_size(n)
        val = codeflash_output  # 249μs -> 223μs (11.7% faster)
        prev = val


# -------------------- Property-Based and Miscellaneous --------------------


@pytest.mark.parametrize(
    "n,ffn_dim_multiplier,multiple_of",
    [
        (0, 1, 256),
        (1, 1.5, 128),
        (10, 0.5, 32),
        (100, 2, 64),
        (256, 1, 256),
        (512, 1.2, 512),
        (1000, 0.8, 256),
        (999, 2.5, 512),
    ],
)
def test_parametrized_various(n, ffn_dim_multiplier, multiple_of):
    # All results should be a multiple of multiple_of
    codeflash_output = compute_intermediate_size(n, ffn_dim_multiplier, multiple_of)
    val = codeflash_output  # 8.97μs -> 7.61μs (17.9% faster)


def test_rounding_behavior():
    # Test that the result rounds up, not down
    # n=100, int(8*100/3)=266, *1.1=292.6, int=292, round up to 300 (if multiple_of=10)
    codeflash_output = compute_intermediate_size(
        100, ffn_dim_multiplier=1.1, multiple_of=10
    )  # 1.36μs -> 1.24μs (9.84% faster)
    # n=100, int(8*100/3)=266, *1.01=268.66, int=268, round up to 270
    codeflash_output = compute_intermediate_size(
        100, ffn_dim_multiplier=1.01, multiple_of=10
    )  # 496ns -> 509ns (2.55% slower)


def test_multiple_of_zero():
    # Should raise ZeroDivisionError if multiple_of is zero
    with pytest.raises(ZeroDivisionError):
        compute_intermediate_size(100, multiple_of=0)  # 1.82μs -> 1.60μs (13.6% 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-compute_intermediate_size-misg3p9t and push.