⚡️ Speed up function gradient_descent by 18,724%

[codeflash-ai]

📄 18,724% (187.24x) speedup for gradient_descent in src/numpy_pandas/statistical_functions.py

⏱️ Runtime : 4.93 seconds → 26.2 milliseconds (best of 247 runs)

📝 Explanation and details

The optimized code achieves a dramatic 18723% speedup by replacing nested Python loops with vectorized NumPy operations, which leverage highly optimized C/Fortran implementations under the hood.

Key Optimizations Applied:

1. Vectorized Prediction Computation: Replaced the nested loop that computed predictions element-by-element with X.dot(weights), eliminating ~29 million individual multiplications and additions in favor of a single optimized matrix-vector multiplication.

2. Vectorized Gradient Computation: Replaced the nested loop for gradient calculation with (X.T @ errors) / m, which computes the gradient in one matrix operation instead of iterating through each feature and sample individually.

3. Vectorized Weight Updates: Replaced the element-wise weight update loop with weights -= learning_rate * gradient, updating all weights simultaneously.

Why This Leads to Speedup:

• BLAS/LAPACK Optimization: NumPy's dot products and matrix operations use highly optimized BLAS libraries that exploit CPU vectorization (SIMD instructions), cache locality, and parallel processing capabilities.

• Loop Overhead Elimination: The original code had ~30 million Python loop iterations (based on profiler data), each with significant interpreter overhead. The vectorized version eliminates this entirely.

• Memory Access Patterns: Vectorized operations have better cache locality and memory bandwidth utilization compared to scattered element-wise access patterns.

Performance Characteristics by Test Case:

• Large-scale scenarios show the most dramatic improvements: Tests with 500+ samples and 20+ features see speedups of 30,000-90,000%, as the vectorization benefits scale quadratically with problem size.

• Small edge cases show modest improvements or slight slowdowns: Single samples or very small datasets (like test_edge_single_sample) show 13-16% slowdowns due to vectorization overhead outweighing benefits at tiny scales.

• Medium-sized problems see consistent 50-300x speedups: Most practical use cases with hundreds of samples and multiple features benefit significantly from the vectorized approach.

The optimization is most effective for typical machine learning scenarios with substantial datasets, where the fixed overhead of vectorization is amortized across many operations.

✅ 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

import numpy as np
# imports
import pytest  # used for our unit tests
from src.numpy_pandas.statistical_functions import gradient_descent

# unit tests

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

def test_basic_single_feature_perfect_fit():
    # Simple linear regression: y = 2x
    X = np.array([[1], [2], [3], [4]])
    y = np.array([2, 4, 6, 8])
    codeflash_output = gradient_descent(X, y, learning_rate=0.1, iterations=500); weights = codeflash_output # 1.39ms -> 1.38ms (0.483% faster)

def test_basic_two_features_perfect_fit():
    # Multiple linear regression: y = 1*x1 + 3*x2
    X = np.array([[1, 2], [2, 1], [3, 0], [0, 3]])
    y = np.array([1*1+3*2, 2*1+3*1, 3*1+3*0, 0*1+3*3])
    codeflash_output = gradient_descent(X, y, learning_rate=0.05, iterations=1000); weights = codeflash_output # 4.73ms -> 2.53ms (87.3% faster)

def test_basic_bias_term():
    # y = 2x + 5, add bias as a feature of 1s
    X = np.array([[1, 1], [2, 1], [3, 1], [4, 1]])
    y = np.array([7, 9, 11, 13])
    codeflash_output = gradient_descent(X, y, learning_rate=0.1, iterations=1000); weights = codeflash_output # 4.74ms -> 2.52ms (88.0% faster)

def test_basic_zero_weights_no_learning():
    # If learning rate is zero, weights should remain zero
    X = np.array([[1, 2], [3, 4]])
    y = np.array([5, 6])
    codeflash_output = gradient_descent(X, y, learning_rate=0.0, iterations=10); weights = codeflash_output # 32.5μs -> 28.1μs (15.7% faster)

def test_basic_no_iterations():
    # If iterations is zero, weights should remain zero
    X = np.array([[1, 2], [3, 4]])
    y = np.array([5, 6])
    codeflash_output = gradient_descent(X, y, learning_rate=0.1, iterations=0); weights = codeflash_output # 792ns -> 542ns (46.1% faster)

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


def test_edge_single_sample():
    # Only one sample, should still work
    X = np.array([[2, 3]])
    y = np.array([13])  # e.g., y = 2*2 + 3*3 = 13
    codeflash_output = gradient_descent(X, y, learning_rate=0.1, iterations=100); weights = codeflash_output # 214μs -> 249μs (13.9% slower)

def test_edge_single_feature():
    # Only one feature, multiple samples
    X = np.array([[1], [2], [3]])
    y = np.array([2, 4, 6])
    codeflash_output = gradient_descent(X, y, learning_rate=0.1, iterations=500); weights = codeflash_output # 1.15ms -> 1.38ms (16.9% slower)

def test_edge_zero_features():
    # Zero features (n=0)
    X = np.empty((3, 0))
    y = np.array([1, 2, 3])
    codeflash_output = gradient_descent(X, y, learning_rate=0.1, iterations=5); weights = codeflash_output # 5.50μs -> 13.7μs (59.8% slower)

def test_edge_zero_targets():
    # All targets are zero, weights should converge to zero
    X = np.array([[1, 2], [3, 4]])
    y = np.array([0, 0])
    codeflash_output = gradient_descent(X, y, learning_rate=0.1, iterations=100); weights = codeflash_output # 302μs -> 254μs (18.6% faster)

def test_edge_large_learning_rate_diverges():
    # Large learning rate should cause divergence (weights become large)
    X = np.array([[1], [2]])
    y = np.array([2, 4])
    codeflash_output = gradient_descent(X, y, learning_rate=1.0, iterations=50); weights = codeflash_output # 94.5μs -> 141μs (33.2% slower)

def test_edge_negative_learning_rate():
    # Negative learning rate should move weights in the wrong direction
    X = np.array([[1], [2]])
    y = np.array([2, 4])
    codeflash_output = gradient_descent(X, y, learning_rate=-0.1, iterations=100); weights = codeflash_output # 186μs -> 277μs (32.8% slower)


def test_edge_nan_inputs():
    # X or y contains NaN
    X = np.array([[1, np.nan], [3, 4]])
    y = np.array([1, 2])
    codeflash_output = gradient_descent(X, y, learning_rate=0.1, iterations=10); weights = codeflash_output # 29.5μs -> 25.0μs (18.2% faster)

def test_edge_inf_inputs():
    # X or y contains Inf
    X = np.array([[1, np.inf], [3, 4]])
    y = np.array([1, 2])
    codeflash_output = gradient_descent(X, y, learning_rate=0.1, iterations=10); weights = codeflash_output # 33.4μs -> 26.6μs (25.5% faster)

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

def test_large_scale_many_samples():
    # 1000 samples, 3 features, y = 2*x1 + 3*x2 + 4*x3
    np.random.seed(0)
    X = np.random.rand(1000, 3)
    true_weights = np.array([2, 3, 4])
    y = X @ true_weights
    codeflash_output = gradient_descent(X, y, learning_rate=0.1, iterations=500); weights = codeflash_output # 520ms -> 1.59ms (32555% faster)

def test_large_scale_many_features():
    # 10 samples, 100 features, y = sum(x_i)
    np.random.seed(1)
    X = np.random.rand(10, 100)
    true_weights = np.ones(100)
    y = X @ true_weights
    codeflash_output = gradient_descent(X, y, learning_rate=0.01, iterations=1000); weights = codeflash_output # 357ms -> 2.22ms (15987% faster)

def test_large_scale_random_noise():
    # 500 samples, 5 features, y = 1.5*x1 - 2.5*x2 + noise
    np.random.seed(2)
    X = np.random.randn(500, 5)
    true_weights = np.array([1.5, -2.5, 0, 0, 0])
    y = X @ true_weights + np.random.normal(0, 0.1, size=500)
    codeflash_output = gradient_descent(X, y, learning_rate=0.05, iterations=500); weights = codeflash_output # 426ms -> 1.45ms (29292% faster)

def test_large_scale_performance():
    # 1000 samples, 10 features, random data, test that it runs in reasonable time
    np.random.seed(3)
    X = np.random.rand(1000, 10)
    y = np.random.rand(1000)
    codeflash_output = gradient_descent(X, y, learning_rate=0.01, iterations=100); weights = codeflash_output # 337ms -> 498μs (67589% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import math

import numpy as np
# imports
import pytest  # used for our unit tests
from src.numpy_pandas.statistical_functions import gradient_descent

# unit tests

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

def test_single_feature_perfect_fit():
    # Single feature, line passes through all points exactly (y = 2x)
    X = np.array([[1], [2], [3]])
    y = np.array([2, 4, 6])
    codeflash_output = gradient_descent(X, y, learning_rate=0.1, iterations=500); w = codeflash_output # 1.15ms -> 1.38ms (16.4% slower)

def test_multiple_features_perfect_fit():
    # Two features, y = 1*x1 + 2*x2
    X = np.array([[1, 2], [2, 1], [3, 0]])
    y = np.array([1*1+2*2, 2*1+1*2, 3*1+0*2])  # [5, 4, 3]
    codeflash_output = gradient_descent(X, y, learning_rate=0.05, iterations=1000); w = codeflash_output # 3.86ms -> 2.51ms (54.0% faster)

def test_zero_learning_rate():
    # Learning rate 0 should not update weights from zero
    X = np.array([[1, 2], [2, 3]])
    y = np.array([1, 2])
    codeflash_output = gradient_descent(X, y, learning_rate=0.0, iterations=100); w = codeflash_output # 303μs -> 253μs (20.0% faster)

def test_zero_iterations():
    # Zero iterations should return initial weights (zeros)
    X = np.array([[1, 2], [2, 3]])
    y = np.array([1, 2])
    codeflash_output = gradient_descent(X, y, learning_rate=0.1, iterations=0); w = codeflash_output # 750ns -> 583ns (28.6% faster)

def test_negative_learning_rate():
    # Negative learning rate should diverge (weights go in wrong direction)
    X = np.array([[1], [2]])
    y = np.array([2, 4])
    codeflash_output = gradient_descent(X, y, learning_rate=-0.1, iterations=100); w = codeflash_output # 186μs -> 279μs (33.2% slower)

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

def test_empty_X_y():
    # X and y are empty arrays
    X = np.empty((0, 2))
    y = np.empty((0,))
    codeflash_output = gradient_descent(X, y, learning_rate=0.1, iterations=10); w = codeflash_output # 30.6μs -> 29.5μs (3.96% faster)

def test_single_sample_multiple_features():
    # One sample, multiple features
    X = np.array([[3, 4, 5]])
    y = np.array([26])  # 2*3 + 3*4 + 1*5 = 6+12+5=23, but let's see what happens
    codeflash_output = gradient_descent(X, y, learning_rate=0.01, iterations=100); w = codeflash_output # 282μs -> 248μs (13.4% faster)
    # With only one sample, weights should fit exactly if possible
    pred = np.dot(X[0], w)

def test_single_feature_constant_y():
    # All y values are the same, should fit w=0 if X is all zeros
    X = np.zeros((5, 1))
    y = np.ones(5)
    codeflash_output = gradient_descent(X, y, learning_rate=0.1, iterations=100); w = codeflash_output # 273μs -> 226μs (20.4% faster)

def test_non_square_X():
    # More features than samples
    X = np.array([[1, 2, 3], [4, 5, 6]])
    y = np.array([14, 32])  # [1*1+2*2+3*3, 4*1+5*2+6*3]
    codeflash_output = gradient_descent(X, y, learning_rate=0.01, iterations=500); w = codeflash_output # 2.06ms -> 1.25ms (63.9% faster)
    # Should fit reasonably well
    preds = np.dot(X, w)

def test_zero_X():
    # X is all zeros, y is not zero
    X = np.zeros((3, 2))
    y = np.array([1, 2, 3])
    codeflash_output = gradient_descent(X, y, learning_rate=0.1, iterations=100); w = codeflash_output # 338μs -> 217μs (55.6% faster)

def test_high_learning_rate():
    # High learning rate should cause divergence (weights become very large)
    X = np.array([[1], [2]])
    y = np.array([2, 4])
    codeflash_output = gradient_descent(X, y, learning_rate=1.0, iterations=50); w = codeflash_output # 94.7μs -> 140μs (32.7% slower)

def test_incorrect_shapes():
    # X and y shapes mismatch should raise an error
    X = np.array([[1, 2], [3, 4]])
    y = np.array([1, 2, 3])
    with pytest.raises(ValueError):
        gradient_descent(X, y, learning_rate=0.1, iterations=10) # 7.12μs -> 3.08μs (131% faster)

def test_nan_in_X_y():
    # X or y contains nan, should propagate to weights
    X = np.array([[1, np.nan], [2, 3]])
    y = np.array([1, 2])
    codeflash_output = gradient_descent(X, y, learning_rate=0.1, iterations=10); w = codeflash_output # 28.6μs -> 24.7μs (16.0% faster)

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

def test_large_dataset():
    # Large m and n, random data, should run and produce weights of correct shape
    m, n = 500, 20
    np.random.seed(42)
    X = np.random.randn(m, n)
    true_w = np.random.randn(n)
    y = X @ true_w + np.random.randn(m) * 0.1  # small noise
    codeflash_output = gradient_descent(X, y, learning_rate=0.01, iterations=500); w = codeflash_output # 1.68s -> 2.09ms (80470% faster)
    # Weights should be close to true_w
    diff = np.linalg.norm(w - true_w)

def test_large_number_of_features():
    # Many features, few samples
    m, n = 10, 500
    np.random.seed(123)
    X = np.random.randn(m, n)
    true_w = np.random.randn(n)
    y = X @ true_w
    codeflash_output = gradient_descent(X, y, learning_rate=0.01, iterations=300); w = codeflash_output # 541ms -> 1.16ms (46684% faster)

def test_large_number_of_samples():
    # Many samples, few features
    m, n = 1000, 2
    np.random.seed(456)
    X = np.random.randn(m, n)
    true_w = np.array([2.0, -1.0])
    y = X @ true_w + np.random.randn(m) * 0.01
    codeflash_output = gradient_descent(X, y, learning_rate=0.05, iterations=300); w = codeflash_output # 208ms -> 914μs (22720% faster)

def test_performance_large_scale():
    # Test that function runs in reasonable time for large but not huge data
    m, n = 500, 50
    np.random.seed(789)
    X = np.random.randn(m, n)
    y = np.random.randn(m)
    codeflash_output = gradient_descent(X, y, learning_rate=0.01, iterations=100); w = codeflash_output # 829ms -> 854μs (96914% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from src.numpy_pandas.statistical_functions import gradient_descent

To edit these changes git checkout codeflash/optimize-gradient_descent-mdp8j1mz and push.