⚡️ Speed up function standard_logei by 12%

[codeflash-ai]

📄 12% (0.12x) speedup for standard_logei in optuna/_gp/acqf.py

⏱️ Runtime : 3.12 milliseconds → 2.79 milliseconds (best of 130 runs)

📝 Explanation and details

The optimized code achieves an 11% speedup through three key optimizations:

1. Eliminated Redundant Computations
The original code computed intermediate values like 0.5 * z and -_SQRT_HALF * z multiple times within the same expression. The optimized version pre-computes these as z_half, minus_z_half_z, and sqrt_half_z, avoiding duplicate tensor operations.

2. More Efficient Condition Checking
Changed from z[(small := z < -25)]).numel() to (small := z < -25).any(). The .any() method is faster for boolean masks as it can short-circuit and doesn't need to count elements, just check if any exist.

3. Cleaner Intermediate Variable Management
The optimized code separates the complex chained operations into clear intermediate steps (erfc_val, exp_val, main_term), which helps PyTorch's optimizer and makes the computation flow more explicit.

Performance Impact by Test Case:

• Best gains (20-35% faster): Small to medium tensors with mostly positive/mixed values where the first condition dominates
• Modest gains (8-15% faster): Large tensors and high-dimensional cases
• Slight slowdown (7% slower): Cases triggering the second condition (z < -25) due to additional variable assignments, but this is rare in practice

The optimizations are particularly effective for the common case where most z values are > -25, making the pre-computed intermediate values highly beneficial.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

import math
from typing import TYPE_CHECKING

# imports
import pytest  # used for our unit tests
from optuna._gp.acqf import standard_logei
from optuna._imports import _LazyImport

if TYPE_CHECKING:
    import torch
else:
    import torch

_SQRT_HALF = math.sqrt(0.5)
_INV_SQRT_2PI = 1 / math.sqrt(2 * math.pi)
_SQRT_HALF_PI = math.sqrt(0.5 * math.pi)
_LOG_SQRT_2PI = math.log(math.sqrt(2 * math.pi))
from optuna._gp.acqf import standard_logei

# unit tests

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


def test_standard_logei_scalar_positive():
    # Test for z=1.0, compare to manual calculation
    import torch
    z = torch.tensor(1.0)
    # ei(z) = z * cdf(z) + pdf(z)
    cdf = 0.5 * torch.special.erfc(-_SQRT_HALF * z)
    pdf = math.exp(-0.5 * z.item() ** 2) * _INV_SQRT_2PI
    ei = z.item() * cdf.item() + pdf
    expected = math.log(ei)
    codeflash_output = standard_logei(z); result = codeflash_output

def test_standard_logei_scalar_negative():
    # Test for z=-1.0, compare to manual calculation
    import torch
    z = torch.tensor(-1.0)
    cdf = 0.5 * torch.special.erfc(-_SQRT_HALF * z)
    pdf = math.exp(-0.5 * z.item() ** 2) * _INV_SQRT_2PI
    ei = z.item() * cdf.item() + pdf
    expected = math.log(ei)
    codeflash_output = standard_logei(z); result = codeflash_output

def test_standard_logei_1d_tensor():
    # Test for 1D tensor input
    import torch
    z = torch.tensor([0.0, 1.0, -1.0])
    codeflash_output = standard_logei(z); result = codeflash_output
    # Compare each element to scalar calculation
    for i, val in enumerate(z):
        cdf = 0.5 * torch.special.erfc(-_SQRT_HALF * val)
        pdf = math.exp(-0.5 * val.item() ** 2) * _INV_SQRT_2PI
        ei = val.item() * cdf.item() + pdf
        expected = math.log(ei)

def test_standard_logei_2d_tensor():
    # Test for 2D tensor input
    import torch
    z = torch.tensor([[0.0, 2.0], [-2.0, 1.0]])
    codeflash_output = standard_logei(z); result = codeflash_output
    # Compare each element to scalar calculation
    for i in range(z.shape[0]):
        for j in range(z.shape[1]):
            val = z[i, j]
            cdf = 0.5 * torch.special.erfc(-_SQRT_HALF * val)
            pdf = math.exp(-0.5 * val.item() ** 2) * _INV_SQRT_2PI
            ei = val.item() * cdf.item() + pdf
            expected = math.log(ei)

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

def test_standard_logei_extremely_large_positive():
    # Test for very large positive z (should be dominated by z * cdf(z))
    import torch
    z = torch.tensor(100.0)
    codeflash_output = standard_logei(z); result = codeflash_output
    # For large z, cdf(z) ~ 1, pdf(z) ~ 0, so ei(z) ~ z
    expected = math.log(z.item())


def test_standard_logei_switch_condition_boundary():
    # Test at the boundary between first and second condition (z=-25)
    import torch
    z = torch.tensor(-25.0)
    codeflash_output = standard_logei(z); result = codeflash_output
    # Should use first condition
    cdf = 0.5 * torch.special.erfc(-_SQRT_HALF * z)
    pdf = math.exp(-0.5 * z.item() ** 2) * _INV_SQRT_2PI
    ei = z.item() * cdf.item() + pdf
    expected = math.log(ei)


def test_standard_logei_nan_input():
    # Test for NaN input, should return NaN
    import torch
    z = torch.tensor(float('nan'))
    codeflash_output = standard_logei(z); result = codeflash_output

def test_standard_logei_inf_input():
    # Test for +inf input, should return +inf
    import torch
    z = torch.tensor(float('inf'))
    codeflash_output = standard_logei(z); result = codeflash_output

def test_standard_logei_neg_inf_input():
    # Test for -inf input, should return -inf
    import torch
    z = torch.tensor(float('-inf'))
    codeflash_output = standard_logei(z); result = codeflash_output


def test_standard_logei_large_tensor_all_positive():
    # Test for large tensor, all positive values
    import torch
    z = torch.linspace(0, 10, steps=1000)
    codeflash_output = standard_logei(z); result = codeflash_output
    # Check monotonicity: for positive z, logei(z) should increase
    for i in range(1, len(result)):
        pass

def test_standard_logei_large_tensor_all_negative():
    # Test for large tensor, all negative values
    import torch
    z = torch.linspace(-100, -1, steps=1000)
    codeflash_output = standard_logei(z); result = codeflash_output
    # For very negative z, logei(z) should decrease
    for i in range(1, len(result)):
        pass


def test_standard_logei_large_tensor_nan_inf():
    # Test for tensor with nan, inf, -inf scattered
    import torch
    z = torch.tensor([float('nan'), float('inf'), float('-inf')] + [0.0]*997)
    codeflash_output = standard_logei(z); result = codeflash_output
    # The rest should be log(pdf(0))
    expected = math.log(_INV_SQRT_2PI)
    for i in range(3, 1000):
        pass

def test_standard_logei_broadcasting():
    # Test broadcasting behavior (scalar z with tensor)
    import torch
    z = torch.tensor([1.0])
    z_large = z.expand(1000)
    codeflash_output = standard_logei(z_large); result = codeflash_output
    # All values should be identical
    cdf = 0.5 * torch.special.erfc(-_SQRT_HALF * z)
    pdf = math.exp(-0.5 * z.item() ** 2) * _INV_SQRT_2PI
    ei = z.item() * cdf.item() + pdf
    expected = math.log(ei)
    for i in range(1000):
        pass

def test_standard_logei_gradient():
    # Test that the function is differentiable for typical values
    import torch
    z = torch.tensor([0.5, -5.0, 10.0], requires_grad=True)
    codeflash_output = standard_logei(z); result = codeflash_output
    result.sum().backward()
    for grad in z.grad:
        pass

def test_standard_logei_dtype_preservation():
    # Test that output dtype matches input dtype
    import torch
    z = torch.tensor([0.0, 1.0, -1.0], dtype=torch.float64)
    codeflash_output = standard_logei(z); result = codeflash_output

def test_standard_logei_device_preservation():
    # Test that output device matches input device (CPU)
    import torch
    z = torch.tensor([0.0, 1.0, -1.0])
    codeflash_output = standard_logei(z); result = codeflash_output

# If CUDA is available, test for GPU device
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")

#------------------------------------------------
from __future__ import annotations

import math
from typing import TYPE_CHECKING

# imports
import pytest  # used for our unit tests
# unit tests
import torch
from optuna._gp.acqf import standard_logei
from optuna._imports import _LazyImport

if TYPE_CHECKING:
    import torch
else:
    import torch

_SQRT_HALF = math.sqrt(0.5)
_INV_SQRT_2PI = 1 / math.sqrt(2 * math.pi)
_SQRT_HALF_PI = math.sqrt(0.5 * math.pi)
_LOG_SQRT_2PI = math.log(math.sqrt(2 * math.pi))
from optuna._gp.acqf import standard_logei

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

def test_standard_logei_zero():
    # Test at z = 0 (should be log(E[max(0, x)]), x~N(0,1))
    z = torch.tensor([0.0])
    codeflash_output = standard_logei(z); result = codeflash_output # 87.0μs -> 74.1μs (17.4% faster)
    # Calculate expected: E[max(0, x)] = 1/sqrt(2*pi) ~ 0.39894228, so log(expected)
    expected = math.log(1 / math.sqrt(2 * math.pi))

def test_standard_logei_positive():
    # Test at positive z values
    z = torch.tensor([1.0, 2.0, 3.0])
    codeflash_output = standard_logei(z); result = codeflash_output # 68.2μs -> 56.4μs (20.9% faster)
    # Compare to manual computation: log(z * cdf(z) + pdf(z))
    for i, zi in enumerate(z):
        cdf = 0.5 * (1 + math.erf(float(zi) / math.sqrt(2)))
        pdf = math.exp(-0.5 * float(zi)**2) / math.sqrt(2 * math.pi)
        expected = math.log(float(zi) * cdf + pdf)

def test_standard_logei_negative():
    # Test at negative z values
    z = torch.tensor([-1.0, -2.0, -3.0])
    codeflash_output = standard_logei(z); result = codeflash_output # 70.3μs -> 55.6μs (26.5% faster)
    for i, zi in enumerate(z):
        cdf = 0.5 * (1 + math.erf(float(zi) / math.sqrt(2)))
        pdf = math.exp(-0.5 * float(zi)**2) / math.sqrt(2 * math.pi)
        expected = math.log(float(zi) * cdf + pdf)

def test_standard_logei_mixed_batch():
    # Test a batch of mixed z values
    z = torch.tensor([-2.0, 0.0, 2.0])
    codeflash_output = standard_logei(z); result = codeflash_output # 67.0μs -> 54.1μs (23.8% faster)
    for i, zi in enumerate(z):
        cdf = 0.5 * (1 + math.erf(float(zi) / math.sqrt(2)))
        pdf = math.exp(-0.5 * float(zi)**2) / math.sqrt(2 * math.pi)
        expected = math.log(float(zi) * cdf + pdf)

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

def test_standard_logei_large_positive():
    # For large positive z, E[max(0, x+z)] ~ z, so log(E) ~ log(z)
    z = torch.tensor([10.0, 20.0, 50.0])
    codeflash_output = standard_logei(z); result = codeflash_output # 69.1μs -> 55.1μs (25.4% faster)
    for i, zi in enumerate(z):
        expected = math.log(float(zi))


def test_standard_logei_extremely_small_z():
    # Test at z near zero (should be smooth, no NaN/Inf)
    z = torch.tensor([1e-10, -1e-10])
    codeflash_output = standard_logei(z); result = codeflash_output # 92.5μs -> 81.3μs (13.7% faster)
    for val in result:
        pass

def test_standard_logei_extremely_large_negative_switch():
    # Test at z just below -25 to ensure second branch is used
    z = torch.tensor([-25.1, -30.0, -100.0])
    codeflash_output = standard_logei(z); result = codeflash_output # 128μs -> 137μs (7.01% slower)
    # Should not be nan or inf
    for val in result:
        pass

def test_standard_logei_extremely_large_positive():
    # Test at z just above 25 (should be in first branch, but still finite)
    z = torch.tensor([25.1, 30.0, 100.0])
    codeflash_output = standard_logei(z); result = codeflash_output # 66.6μs -> 61.3μs (8.66% faster)
    for val in result:
        pass

def test_standard_logei_shape_preservation():
    # Test that input shape is preserved
    z = torch.randn(5, 3, 2)
    codeflash_output = standard_logei(z); result = codeflash_output # 76.1μs -> 57.0μs (33.4% faster)

def test_standard_logei_dtype_preservation():
    # Test that dtype is preserved (float32, float64)
    z32 = torch.randn(4, dtype=torch.float32)
    z64 = torch.randn(4, dtype=torch.float64)
    codeflash_output = standard_logei(z32).dtype # 62.8μs -> 51.9μs (21.1% faster)
    codeflash_output = standard_logei(z64).dtype # 32.7μs -> 29.2μs (11.8% faster)

def test_standard_logei_requires_grad():
    # Test that gradients can be computed (autograd support)
    z = torch.tensor([0.5, -0.5], requires_grad=True)
    codeflash_output = standard_logei(z); out = codeflash_output # 95.5μs -> 79.0μs (20.9% faster)
    s = out.sum()
    s.backward()

def test_standard_logei_nan_inf_input():
    # Test for NaN and Inf input
    z = torch.tensor([float('nan'), float('inf'), float('-inf')])
    codeflash_output = standard_logei(z); result = codeflash_output # 126μs -> 137μs (7.45% slower)

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

def test_standard_logei_large_batch():
    # Test on a large batch (max 1000 elements)
    z = torch.linspace(-30, 30, 1000)
    codeflash_output = standard_logei(z); result = codeflash_output # 147μs -> 159μs (7.18% slower)

def test_standard_logei_large_random_tensor():
    # Test on a large random tensor (max 1000 elements, 2D)
    z = torch.randn(100, 10)
    codeflash_output = standard_logei(z); result = codeflash_output # 77.3μs -> 67.1μs (15.1% faster)

def test_standard_logei_gradient_large_batch():
    # Test autograd on large batch
    z = torch.randn(1000, requires_grad=True)
    codeflash_output = standard_logei(z); out = codeflash_output # 100μs -> 83.5μs (20.0% faster)
    s = out.sum()
    s.backward()

def test_standard_logei_high_dimensional_tensor():
    # Test on a high dimensional tensor (but <100MB)
    z = torch.randn(10, 5, 5, 4)
    codeflash_output = standard_logei(z); result = codeflash_output # 91.8μs -> 73.9μs (24.3% faster)

# ----------- Miscellaneous -----------

def test_standard_logei_empty_tensor():
    # Test on an empty tensor
    z = torch.tensor([])
    codeflash_output = standard_logei(z); result = codeflash_output # 62.5μs -> 47.3μs (32.0% faster)

def test_standard_logei_scalar_input():
    # Test on a scalar input (0-dim tensor)
    z = torch.tensor(1.23)
    codeflash_output = standard_logei(z); result = codeflash_output # 69.9μs -> 53.6μs (30.5% faster)
    # Compare to manual calculation
    cdf = 0.5 * (1 + math.erf(float(z) / math.sqrt(2)))
    pdf = math.exp(-0.5 * float(z)**2) / math.sqrt(2 * math.pi)
    expected = math.log(float(z) * cdf + pdf)

def test_standard_logei_noncontiguous_tensor():
    # Test with non-contiguous tensor
    z = torch.randn(10, 10).t()[::2]
    codeflash_output = standard_logei(z); result = codeflash_output # 74.9μs -> 59.8μs (25.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-standard_logei-mhbf6tc7 and push.