⚡️ Speed up function calculate_shift by 8%

[codeflash-ai]

📄 8% (0.08x) speedup for calculate_shift in src/diffusers/pipelines/ltx/pipeline_ltx.py

⏱️ Runtime : 639 microseconds → 594 microseconds (best of 457 runs)

📝 Explanation and details

Here is an optimized rewrite of your program for runtime and memory efficiency. (There was little room for micro-optimization in such a small, purely-arithmetic function, but every bit helps!)

• Inline calculation expressions when possible (minimize assignments).
• Use constant folding when possible with default arguments (so m and b are computed only once for default parameters).
• Use __slots__ in the helper (if a callable object is used).
• Remove redundant assignments.

Note: This function is already very efficient, so changes are incremental and technical.

Summary of optimizations:

• Precomputes constants for default parameters and uses an optimized code path in the common case.
• Uses less memory by removing the mu assignment and unneeded intermediate variables.
• Only creates variables for the general case, avoiding unnecessary calculation in the default/common case.

Function signature and return values remain exactly the same.

✅ Correctness verification report:

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

🌀 Generated Regression Tests Details

from math import isclose

# imports
import pytest  # used for our unit tests
from src.diffusers.pipelines.ltx.pipeline_ltx import calculate_shift
from src.diffusers.utils import is_torch_xla_available

# unit tests

# ----------------- BASIC TEST CASES -----------------

def test_exactly_base_seq_len():
    # At base_seq_len, should return base_shift
    codeflash_output = calculate_shift(256); result = codeflash_output

def test_exactly_max_seq_len():
    # At max_seq_len, should return max_shift
    codeflash_output = calculate_shift(4096); result = codeflash_output

def test_midpoint_seq_len():
    # At midpoint between base and max, should be halfway between shifts
    mid_seq = (256 + 4096) // 2
    expected = (0.5 + 1.15) / 2
    codeflash_output = calculate_shift(mid_seq); result = codeflash_output

def test_typical_in_between():
    # Typical value between base and max
    seq = 1024
    # Manually compute expected
    m = (1.15 - 0.5) / (4096 - 256)
    b = 0.5 - m * 256
    expected = seq * m + b
    codeflash_output = calculate_shift(seq); result = codeflash_output

def test_non_default_parameters():
    # Use custom base/max values
    codeflash_output = calculate_shift(
        image_seq_len=1000,
        base_seq_len=100,
        max_seq_len=2000,
        base_shift=1.0,
        max_shift=2.0,
    ); result = codeflash_output
    m = (2.0 - 1.0) / (2000 - 100)
    b = 1.0 - m * 100
    expected = 1000 * m + b

# ----------------- EDGE TEST CASES -----------------

def test_below_base_seq_len():
    # image_seq_len below base_seq_len (extrapolation)
    seq = 0
    m = (1.15 - 0.5) / (4096 - 256)
    b = 0.5 - m * 256
    expected = seq * m + b
    codeflash_output = calculate_shift(seq); result = codeflash_output

def test_above_max_seq_len():
    # image_seq_len above max_seq_len (extrapolation)
    seq = 5000
    m = (1.15 - 0.5) / (4096 - 256)
    b = 0.5 - m * 256
    expected = seq * m + b
    codeflash_output = calculate_shift(seq); result = codeflash_output

def test_negative_seq_len():
    # Negative image_seq_len (nonsensical, but should still compute)
    seq = -100
    m = (1.15 - 0.5) / (4096 - 256)
    b = 0.5 - m * 256
    expected = seq * m + b
    codeflash_output = calculate_shift(seq); result = codeflash_output

def test_base_equals_max_seq_len():
    # base_seq_len == max_seq_len should raise ZeroDivisionError
    with pytest.raises(ZeroDivisionError):
        calculate_shift(100, base_seq_len=100, max_seq_len=100)

def test_base_greater_than_max_seq_len():
    # base_seq_len > max_seq_len should still compute (negative slope)
    codeflash_output = calculate_shift(50, base_seq_len=100, max_seq_len=10, base_shift=1.0, max_shift=0.0); result = codeflash_output
    m = (0.0 - 1.0) / (10 - 100)
    b = 1.0 - m * 100
    expected = 50 * m + b

def test_float_seq_len():
    # image_seq_len as float
    seq = 512.5
    m = (1.15 - 0.5) / (4096 - 256)
    b = 0.5 - m * 256
    expected = seq * m + b
    codeflash_output = calculate_shift(seq); result = codeflash_output

def test_large_float_shifts():
    # Very large shift values
    codeflash_output = calculate_shift(3000, base_seq_len=100, max_seq_len=2000, base_shift=1e6, max_shift=1e9); result = codeflash_output
    m = (1e9 - 1e6) / (2000 - 100)
    b = 1e6 - m * 100
    expected = 3000 * m + b

def test_all_zero_parameters():
    # All zero parameters
    codeflash_output = calculate_shift(0, base_seq_len=0, max_seq_len=1, base_shift=0.0, max_shift=0.0); result = codeflash_output

def test_base_shift_equals_max_shift():
    # base_shift == max_shift: always returns base_shift
    for seq in [0, 100, 256, 4096, 10000]:
        codeflash_output = calculate_shift(seq, base_seq_len=256, max_seq_len=4096, base_shift=0.7, max_shift=0.7); result = codeflash_output

# ----------------- LARGE SCALE TEST CASES -----------------

def test_large_range_of_seq_lens():
    # Test over a large range of image_seq_len values (0 to 999)
    m = (1.15 - 0.5) / (4096 - 256)
    b = 0.5 - m * 256
    for seq in range(0, 1000):
        expected = seq * m + b
        codeflash_output = calculate_shift(seq); result = codeflash_output

def test_many_parameter_combinations():
    # Test many combinations of base/max seq_len and shift, within reasonable bounds
    for base_seq_len in range(10, 110, 20):
        for max_seq_len in range(base_seq_len+1, base_seq_len+101, 30):
            for base_shift in [0.0, 0.5, 1.0]:
                for max_shift in [base_shift+0.1, base_shift+1.0]:
                    for seq in [base_seq_len, max_seq_len, (base_seq_len+max_seq_len)//2]:
                        m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
                        b = base_shift - m * base_seq_len
                        expected = seq * m + b
                        codeflash_output = calculate_shift(seq, base_seq_len, max_seq_len, base_shift, max_shift); result = codeflash_output

def test_performance_large_inputs():
    # Test with large but manageable parameters (no more than 1000 elements)
    base_seq_len = 1
    max_seq_len = 1000
    base_shift = 0.0
    max_shift = 10.0
    # Test at various points
    for seq in [1, 500, 999, 1000]:
        m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
        b = base_shift - m * base_seq_len
        expected = seq * m + b
        codeflash_output = calculate_shift(seq, base_seq_len, max_seq_len, base_shift, max_shift); result = codeflash_output

def test_stress_many_calls():
    # Stress test: call function 1000 times with increasing seq_len
    base_seq_len = 10
    max_seq_len = 1010
    base_shift = 2.0
    max_shift = 5.0
    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
    b = base_shift - m * base_seq_len
    for seq in range(10, 1010):
        expected = seq * m + b
        codeflash_output = calculate_shift(seq, base_seq_len, max_seq_len, base_shift, max_shift); result = codeflash_output
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from math import isclose

# imports
import pytest  # used for our unit tests
from src.diffusers.pipelines.ltx.pipeline_ltx import calculate_shift

# unit tests

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

def test_calculate_shift_default_base_seq_len():
    # Test at the default base_seq_len: should return base_shift
    codeflash_output = calculate_shift(256); result = codeflash_output

def test_calculate_shift_default_max_seq_len():
    # Test at the default max_seq_len: should return max_shift
    codeflash_output = calculate_shift(4096); result = codeflash_output

def test_calculate_shift_middle_value():
    # Test at the midpoint between base_seq_len and max_seq_len
    mid_seq_len = (256 + 4096) // 2
    # Calculate expected value
    m = (1.15 - 0.5) / (4096 - 256)
    b = 0.5 - m * 256
    expected = mid_seq_len * m + b
    codeflash_output = calculate_shift(mid_seq_len); result = codeflash_output

def test_calculate_shift_custom_base_and_max():
    # Test with custom base_seq_len and max_seq_len
    codeflash_output = calculate_shift(1000, base_seq_len=1000, max_seq_len=2000, base_shift=0.2, max_shift=0.8); result = codeflash_output
    # At image_seq_len=2000, should get max_shift
    codeflash_output = calculate_shift(2000, base_seq_len=1000, max_seq_len=2000, base_shift=0.2, max_shift=0.8); result2 = codeflash_output

def test_calculate_shift_negative_shift():
    # Test with negative shift values
    codeflash_output = calculate_shift(300, base_seq_len=200, max_seq_len=400, base_shift=-1.0, max_shift=1.0); result = codeflash_output
    m = (1.0 - (-1.0)) / (400 - 200)
    b = -1.0 - m * 200
    expected = 300 * m + b

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

def test_calculate_shift_image_seq_len_below_base():
    # image_seq_len less than base_seq_len (extrapolation)
    codeflash_output = calculate_shift(100); result = codeflash_output
    m = (1.15 - 0.5) / (4096 - 256)
    b = 0.5 - m * 256
    expected = 100 * m + b

def test_calculate_shift_image_seq_len_above_max():
    # image_seq_len greater than max_seq_len (extrapolation)
    codeflash_output = calculate_shift(5000); result = codeflash_output
    m = (1.15 - 0.5) / (4096 - 256)
    b = 0.5 - m * 256
    expected = 5000 * m + b

def test_calculate_shift_base_equals_max_seq_len():
    # base_seq_len == max_seq_len should raise ZeroDivisionError
    with pytest.raises(ZeroDivisionError):
        calculate_shift(100, base_seq_len=100, max_seq_len=100, base_shift=0.5, max_shift=1.5)

def test_calculate_shift_base_shift_equals_max_shift():
    # base_shift == max_shift should produce constant output for any image_seq_len
    codeflash_output = calculate_shift(256, base_seq_len=256, max_seq_len=4096, base_shift=1.0, max_shift=1.0); result1 = codeflash_output
    codeflash_output = calculate_shift(4096, base_seq_len=256, max_seq_len=4096, base_shift=1.0, max_shift=1.0); result2 = codeflash_output
    codeflash_output = calculate_shift(1000, base_seq_len=256, max_seq_len=4096, base_shift=1.0, max_shift=1.0); result3 = codeflash_output

def test_calculate_shift_negative_seq_lens():
    # Negative image_seq_len should be handled mathematically
    codeflash_output = calculate_shift(-10); result = codeflash_output
    m = (1.15 - 0.5) / (4096 - 256)
    b = 0.5 - m * 256
    expected = -10 * m + b

def test_calculate_shift_float_inputs():
    # Test with float image_seq_len
    codeflash_output = calculate_shift(512.5); result = codeflash_output
    m = (1.15 - 0.5) / (4096 - 256)
    b = 0.5 - m * 256
    expected = 512.5 * m + b

def test_calculate_shift_float_base_and_max_seq_len():
    # Test with float base_seq_len and max_seq_len
    codeflash_output = calculate_shift(300.0, base_seq_len=200.0, max_seq_len=400.0, base_shift=0.0, max_shift=1.0); result = codeflash_output
    m = (1.0 - 0.0) / (400.0 - 200.0)
    b = 0.0 - m * 200.0
    expected = 300.0 * m + b

def test_calculate_shift_large_negative_values():
    # Large negative image_seq_len
    codeflash_output = calculate_shift(-10000); result = codeflash_output
    m = (1.15 - 0.5) / (4096 - 256)
    b = 0.5 - m * 256
    expected = -10000 * m + b

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

def test_calculate_shift_large_range():
    # Test a large range of image_seq_len values for monotonicity and correctness
    m = (1.15 - 0.5) / (4096 - 256)
    b = 0.5 - m * 256
    codeflash_output = calculate_shift(0); prev_result = codeflash_output
    for i in range(1, 1000):  # up to 999
        codeflash_output = calculate_shift(i); result = codeflash_output
        expected = i * m + b
        prev_result = result

def test_calculate_shift_large_custom_params():
    # Test with large custom parameters and large image_seq_len values
    base_seq_len = 10
    max_seq_len = 1000
    base_shift = 0.0
    max_shift = 10.0
    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
    b = base_shift - m * base_seq_len
    for i in range(base_seq_len, max_seq_len + 1, 100):
        codeflash_output = calculate_shift(i, base_seq_len=base_seq_len, max_seq_len=max_seq_len, base_shift=base_shift, max_shift=max_shift); result = codeflash_output
        expected = i * m + b

def test_calculate_shift_performance_large_inputs():
    # Test performance and correctness for image_seq_len near upper bound
    codeflash_output = calculate_shift(999); result = codeflash_output
    m = (1.15 - 0.5) / (4096 - 256)
    b = 0.5 - m * 256
    expected = 999 * m + b

def test_calculate_shift_large_float_inputs():
    # Test with large float image_seq_len values
    for i in range(0, 1000, 100):
        image_seq_len = float(i) + 0.5
        codeflash_output = calculate_shift(image_seq_len); result = codeflash_output
        m = (1.15 - 0.5) / (4096 - 256)
        b = 0.5 - m * 256
        expected = image_seq_len * m + b
# 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-calculate_shift-mbdgxab0 and push.