⚡️ Speed up method IntDistribution.to_internal_repr by 10%

[codeflash-ai]

📄 10% (0.10x) speedup for IntDistribution.to_internal_repr in optuna/distributions.py

⏱️ Runtime : 1.39 milliseconds → 1.27 milliseconds (best of 149 runs)

📝 Explanation and details

The optimized code achieves a 9% speedup by replacing the math.isnan() function call with a direct NaN comparison (internal_repr != internal_repr). This is the only meaningful change between the versions.

Key optimization:

• NaN detection improvement: Changed math.isnan(internal_repr) to internal_repr != internal_repr
  • Eliminates a function call overhead to the math module
  • Uses Python's built-in float comparison behavior where NaN != NaN always returns True
  • Reduces execution time from ~280ns to ~229ns per hit (18% faster for this specific check)

Why this works:
In IEEE 754 floating-point arithmetic (Python's float standard), NaN has the unique property that it never equals itself. The expression x != x is the idiomatic and fastest way to test for NaN in Python, avoiding the overhead of importing and calling math.isnan().

Performance characteristics:

• Best gains (15-35% faster): Valid input test cases where the NaN check executes successfully
• Minimal impact on error cases: Exception handling paths show little to no improvement since they're dominated by exception creation overhead
• Consistent improvement: All test cases with valid numeric inputs benefit from the faster NaN detection

The optimization is particularly effective for high-frequency parameter validation scenarios typical in hyperparameter optimization frameworks like Optuna.

✅ Correctness verification report:

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

⚙️ Existing Unit Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
test_distributions.py::test_int_internal_representation	9.17μs	7.55μs	21.5%✅
test_distributions.py::test_int_internal_representation_error	201μs	203μs	-1.00%⚠️

🌀 Generated Regression Tests and Runtime

import math
import warnings

# imports
import pytest  # used for our unit tests
from optuna.distributions import IntDistribution


# function to test
class BaseDistribution:
    pass
from optuna.distributions import IntDistribution


def _adjust_int_uniform_high(low: int, high: int, step: int) -> int:
    r = high - low
    if r % step != 0:
        old_high = high
        high = r // step * step + low
        warnings.warn(
            f"The distribution is specified by [{low}, {old_high}] and {step=}, but the range is "
            f"not divisible by `step`. It will be replaced with [{low}, {high}]."
        )
    return high

# unit tests

# --- Basic Test Cases ---
def test_basic_positive_int():
    """Test with a normal positive integer value."""
    dist = IntDistribution(low=1, high=10)
    codeflash_output = dist.to_internal_repr(5) # 772ns -> 654ns (18.0% faster)

def test_basic_negative_int():
    """Test with a normal negative integer value (allowed if log=False)."""
    dist = IntDistribution(low=-10, high=10)
    codeflash_output = dist.to_internal_repr(-3) # 755ns -> 655ns (15.3% faster)

def test_basic_zero():
    """Test with zero value."""
    dist = IntDistribution(low=0, high=10)
    codeflash_output = dist.to_internal_repr(0) # 721ns -> 650ns (10.9% faster)

def test_basic_float_castable_string():
    """Test with a string that can be cast to float."""
    dist = IntDistribution(low=1, high=10)
    codeflash_output = dist.to_internal_repr("7") # 1.12μs -> 1.05μs (7.07% faster)

def test_basic_float_castable_float():
    """Test with a float value."""
    dist = IntDistribution(low=1, high=10)
    codeflash_output = dist.to_internal_repr(3.5) # 649ns -> 576ns (12.7% faster)

def test_basic_high_equals_low():
    """Test with low == high, only one value allowed."""
    dist = IntDistribution(low=5, high=5)
    codeflash_output = dist.to_internal_repr(5) # 754ns -> 667ns (13.0% faster)

# --- Edge Test Cases ---
def test_invalid_type_string():
    """Test with a string that cannot be cast to float."""
    dist = IntDistribution(low=1, high=10)
    with pytest.raises(ValueError):
        dist.to_internal_repr("abc") # 2.83μs -> 3.15μs (10.0% slower)

def test_invalid_type_none():
    """Test with None as input."""
    dist = IntDistribution(low=1, high=10)
    with pytest.raises(ValueError):
        dist.to_internal_repr(None) # 2.43μs -> 2.48μs (2.13% slower)

def test_invalid_type_list():
    """Test with a list as input."""
    dist = IntDistribution(low=1, high=10)
    with pytest.raises(ValueError):
        dist.to_internal_repr([1,2,3]) # 3.55μs -> 4.02μs (11.6% slower)

def test_nan_value():
    """Test with NaN value."""
    dist = IntDistribution(low=1, high=10)
    with pytest.raises(ValueError):
        dist.to_internal_repr(float('nan')) # 2.23μs -> 2.11μs (5.73% faster)

def test_log_distribution_positive():
    """Test with log=True and positive value."""
    dist = IntDistribution(low=1, high=10, log=True)
    codeflash_output = dist.to_internal_repr(2) # 973ns -> 822ns (18.4% faster)

def test_log_distribution_zero():
    """Test with log=True and zero value (should fail)."""
    dist = IntDistribution(low=1, high=10, log=True)
    with pytest.raises(ValueError):
        dist.to_internal_repr(0) # 1.52μs -> 1.51μs (0.597% faster)

def test_log_distribution_negative():
    """Test with log=True and negative value (should fail)."""
    dist = IntDistribution(low=1, high=10, log=True)
    with pytest.raises(ValueError):
        dist.to_internal_repr(-5) # 1.69μs -> 1.58μs (6.89% faster)

def test_step_adjustment_warning():
    """Test that step adjustment works and warning is issued."""
    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        dist = IntDistribution(low=1, high=10, step=3)

def test_step_is_zero():
    """Test that step=0 raises ValueError."""
    with pytest.raises(ValueError):
        IntDistribution(low=1, high=10, step=0)

def test_step_is_negative():
    """Test that step<0 raises ValueError."""
    with pytest.raises(ValueError):
        IntDistribution(low=1, high=10, step=-2)

def test_log_and_step_not_one():
    """Test that log=True and step!=1 raises ValueError."""
    with pytest.raises(ValueError):
        IntDistribution(low=1, high=10, log=True, step=2)

def test_low_greater_than_high():
    """Test that low > high raises ValueError."""
    with pytest.raises(ValueError):
        IntDistribution(low=10, high=1)

def test_log_and_low_less_than_one():
    """Test that log=True and low<1 raises ValueError."""
    with pytest.raises(ValueError):
        IntDistribution(low=0, high=10, log=True)

def test_to_internal_repr_outside_range():
    """Test to_internal_repr with value outside the distribution range (allowed)."""
    dist = IntDistribution(low=1, high=10)
    # No range check in to_internal_repr; should succeed
    codeflash_output = dist.to_internal_repr(100) # 1.07μs -> 791ns (34.9% faster)
    codeflash_output = dist.to_internal_repr(-100) # 348ns -> 281ns (23.8% faster)

def test_to_internal_repr_bool():
    """Test with boolean values."""
    dist = IntDistribution(low=0, high=1)
    codeflash_output = dist.to_internal_repr(True) # 713ns -> 633ns (12.6% faster)
    codeflash_output = dist.to_internal_repr(False) # 338ns -> 275ns (22.9% faster)

def test_to_internal_repr_inf():
    """Test with infinity."""
    dist = IntDistribution(low=1, high=10)
    codeflash_output = dist.to_internal_repr(float("inf")) # 665ns -> 554ns (20.0% faster)
    codeflash_output = dist.to_internal_repr(float("-inf")) # 279ns -> 237ns (17.7% faster)

# --- Large Scale Test Cases ---
def test_large_scale_range():
    """Test with a large range of values."""
    dist = IntDistribution(low=1, high=999)
    for i in range(1, 1000, 100):  # sample 10 values
        codeflash_output = dist.to_internal_repr(i) # 2.67μs -> 2.35μs (13.6% faster)

def test_large_scale_string_inputs():
    """Test with many string inputs that can be cast to float."""
    dist = IntDistribution(low=1, high=1000)
    for i in range(1, 1000, 111):  # sample 9 values
        codeflash_output = dist.to_internal_repr(str(i)) # 3.31μs -> 2.90μs (14.2% faster)

def test_large_scale_invalid_strings():
    """Test with many invalid string inputs."""
    dist = IntDistribution(low=1, high=1000)
    invalids = ["foo", "bar", "baz", "qux", "quux", "corge", "grault", "garply", "waldo"]
    for s in invalids:
        with pytest.raises(ValueError):
            dist.to_internal_repr(s)

def test_large_scale_nan_inputs():
    """Test with many NaN inputs."""
    dist = IntDistribution(low=1, high=1000)
    for _ in range(10):
        with pytest.raises(ValueError):
            dist.to_internal_repr(float('nan'))

def test_large_scale_log_distribution():
    """Test log distribution with a large range."""
    dist = IntDistribution(low=1, high=1000, log=True)
    for i in range(1, 1000, 111):  # sample 9 values
        codeflash_output = dist.to_internal_repr(i) # 2.81μs -> 2.46μs (14.3% faster)
    for i in range(-100, 0, 10):  # sample 10 negative values
        with pytest.raises(ValueError):
            dist.to_internal_repr(i)

def test_large_scale_step_adjustment():
    """Test step adjustment on large range."""
    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        dist = IntDistribution(low=1, high=999, step=100)

def test_large_scale_to_internal_repr_bool():
    """Test with many boolean values."""
    dist = IntDistribution(low=0, high=1)
    for _ in range(100):
        codeflash_output = dist.to_internal_repr(True) # 22.2μs -> 19.0μs (16.9% faster)
        codeflash_output = dist.to_internal_repr(False)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from __future__ import annotations

import math
import warnings

# imports
import pytest  # used for our unit tests
from optuna.distributions import IntDistribution


class BaseDistribution:
    pass
from optuna.distributions import IntDistribution


def _adjust_int_uniform_high(low: int, high: int, step: int) -> int:
    r = high - low
    if r % step != 0:
        old_high = high
        high = r // step * step + low
        warnings.warn(
            f"The distribution is specified by [{low}, {old_high}] and {step=}, but the range is "
            f"not divisible by `step`. It will be replaced with [{low}, {high}]."
        )
    return high

# unit tests

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

def test_basic_integer_input():
    # Test with a basic integer value in normal range
    dist = IntDistribution(low=1, high=10)
    codeflash_output = dist.to_internal_repr(5) # 821ns -> 738ns (11.2% faster)

def test_basic_float_input():
    # Test with a float value in normal range
    dist = IntDistribution(low=1, high=10)
    codeflash_output = dist.to_internal_repr(5.5) # 710ns -> 616ns (15.3% faster)

def test_basic_string_castable_input():
    # Test with a string that can be cast to float
    dist = IntDistribution(low=1, high=10)
    codeflash_output = dist.to_internal_repr("7") # 1.24μs -> 1.08μs (14.8% faster)
    codeflash_output = dist.to_internal_repr("7.5") # 516ns -> 488ns (5.74% faster)

def test_basic_log_distribution_positive():
    # Test with log=True and positive value
    dist = IntDistribution(low=1, high=10, log=True)
    codeflash_output = dist.to_internal_repr(2) # 912ns -> 761ns (19.8% faster)
    codeflash_output = dist.to_internal_repr("3") # 572ns -> 591ns (3.21% slower)

def test_basic_step_distribution():
    # Test with step > 1
    dist = IntDistribution(low=0, high=10, step=2)
    codeflash_output = dist.to_internal_repr(4) # 714ns -> 661ns (8.02% faster)

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

def test_edge_low_equals_high():
    # low == high, only one valid value
    dist = IntDistribution(low=5, high=5)
    codeflash_output = dist.to_internal_repr(5) # 707ns -> 586ns (20.6% faster)

def test_edge_low_greater_than_high():
    # low > high should raise ValueError
    with pytest.raises(ValueError):
        IntDistribution(low=10, high=5)

def test_edge_zero_step():
    # step == 0 should raise ValueError
    with pytest.raises(ValueError):
        IntDistribution(low=1, high=10, step=0)

def test_edge_negative_step():
    # step < 0 should raise ValueError
    with pytest.raises(ValueError):
        IntDistribution(low=1, high=10, step=-1)

def test_edge_log_true_and_step_not_one():
    # log=True and step!=1 should raise ValueError
    with pytest.raises(ValueError):
        IntDistribution(low=1, high=10, log=True, step=2)

def test_edge_log_true_and_low_less_than_one():
    # log=True and low<1 should raise ValueError
    with pytest.raises(ValueError):
        IntDistribution(low=0, high=10, log=True)

def test_edge_log_true_and_zero_input():
    # log=True and input <= 0 should raise ValueError
    dist = IntDistribution(low=1, high=10, log=True)
    with pytest.raises(ValueError):
        dist.to_internal_repr(0) # 1.96μs -> 1.60μs (21.9% faster)
    with pytest.raises(ValueError):
        dist.to_internal_repr(-1) # 1.05μs -> 1.02μs (2.15% faster)
    with pytest.raises(ValueError):
        dist.to_internal_repr("0") # 1.05μs -> 1.08μs (2.59% slower)
    with pytest.raises(ValueError):
        dist.to_internal_repr("-2") # 884ns -> 820ns (7.80% faster)

def test_edge_nan_input():
    # input is NaN should raise ValueError
    dist = IntDistribution(low=1, high=10)
    with pytest.raises(ValueError):
        dist.to_internal_repr(float('nan')) # 2.13μs -> 1.99μs (7.30% faster)

def test_edge_invalid_string_input():
    # input is a string that cannot be cast to float
    dist = IntDistribution(low=1, high=10)
    with pytest.raises(ValueError):
        dist.to_internal_repr("abc") # 2.72μs -> 2.89μs (5.91% slower)

def test_edge_none_input():
    # input is None should raise ValueError
    dist = IntDistribution(low=1, high=10)
    with pytest.raises(ValueError):
        dist.to_internal_repr(None) # 2.43μs -> 2.46μs (1.22% slower)

def test_edge_list_input():
    # input is a list should raise ValueError
    dist = IntDistribution(low=1, high=10)
    with pytest.raises(ValueError):
        dist.to_internal_repr([1, 2, 3]) # 3.55μs -> 3.71μs (4.26% slower)

def test_edge_dict_input():
    # input is a dict should raise ValueError
    dist = IntDistribution(low=1, high=10)
    with pytest.raises(ValueError):
        dist.to_internal_repr({'a': 1}) # 3.29μs -> 3.31μs (0.664% slower)

def test_edge_bool_input():
    # input is a boolean, should be cast to float (True->1.0, False->0.0)
    dist = IntDistribution(low=0, high=10)
    codeflash_output = dist.to_internal_repr(True) # 1.03μs -> 781ns (32.1% faster)
    codeflash_output = dist.to_internal_repr(False) # 325ns -> 331ns (1.81% slower)

def test_edge_log_true_and_bool_false_input():
    # log=True and input False (0.0) should raise ValueError
    dist = IntDistribution(low=1, high=10, log=True)
    with pytest.raises(ValueError):
        dist.to_internal_repr(False) # 2.18μs -> 2.00μs (9.31% faster)

def test_edge_high_not_divisible_by_step():
    # high not divisible by step, should adjust high and warn
    with pytest.warns(UserWarning):
        dist = IntDistribution(low=0, high=9, step=2)
    codeflash_output = dist.to_internal_repr(2) # 791ns -> 696ns (13.6% faster)

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

def test_large_scale_many_inputs():
    # Test with a large number of valid integer inputs
    dist = IntDistribution(low=1, high=1000)
    for i in range(1, 1001):
        codeflash_output = dist.to_internal_repr(i) # 203μs -> 174μs (17.1% faster)

def test_large_scale_many_float_inputs():
    # Test with a large number of valid float inputs
    dist = IntDistribution(low=1, high=1000)
    for i in range(1, 1001):
        val = i + 0.5
        codeflash_output = dist.to_internal_repr(val) # 201μs -> 172μs (17.0% faster)

def test_large_scale_many_string_inputs():
    # Test with a large number of string inputs that can be cast to float
    dist = IntDistribution(low=1, high=1000)
    for i in range(1, 1001):
        val = str(i)
        codeflash_output = dist.to_internal_repr(val) # 236μs -> 211μs (12.1% faster)

def test_large_scale_log_distribution():
    # Test log distribution with many valid inputs
    dist = IntDistribution(low=1, high=1000, log=True)
    for i in range(1, 1001):
        codeflash_output = dist.to_internal_repr(i) # 218μs -> 189μs (15.4% faster)

def test_large_scale_invalid_inputs():
    # Test with many invalid inputs (non-castable strings)
    dist = IntDistribution(low=1, high=1000)
    for i in range(100):
        with pytest.raises(ValueError):
            dist.to_internal_repr(f"invalid_{i}")

def test_large_scale_edge_nan_inputs():
    # Test with many nan inputs
    dist = IntDistribution(low=1, high=1000)
    for _ in range(100):
        with pytest.raises(ValueError):
            dist.to_internal_repr(float('nan'))

def test_large_scale_edge_none_inputs():
    # Test with many None inputs
    dist = IntDistribution(low=1, high=1000)
    for _ in range(100):
        with pytest.raises(ValueError):
            dist.to_internal_repr(None)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from optuna.distributions import IntDistribution
import pytest

def test_IntDistribution_to_internal_repr():
    with pytest.raises(ValueError, match='`0`\\ is\\ invalid\\ value\\ for\\ the\\ case\\ log=True\\.'):
        IntDistribution.to_internal_repr(IntDistribution(1, 1, log=True, step=1), 0)

def test_IntDistribution_to_internal_repr_2():
    IntDistribution.to_internal_repr(IntDistribution(1, 1, log=True, step=1), 1)

🔎 Concolic Coverage Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
codeflash_concolic_qluqolhr/tmpez7lxxhx/test_concolic_coverage.py::test_IntDistribution_to_internal_repr	1.82μs	1.68μs	8.89%✅
codeflash_concolic_qluqolhr/tmpez7lxxhx/test_concolic_coverage.py::test_IntDistribution_to_internal_repr_2	880ns	776ns	13.4%✅

To edit these changes git checkout codeflash/optimize-IntDistribution.to_internal_repr-mhbikfx4 and push.