⚡️ Speed up function `_tridiagonal_forward_body_tf` by 72% by codeflash-ai[bot] · Pull Request #1078 · codeflash-ai/codeflash

codeflash-ai · 2026-01-16T08:53:18Z

📄 72% (0.72x) speedup for `_tridiagonal_forward_body_tf` in `code_to_optimize/sample_code.py`

⏱️ Runtime : 407 milliseconds → 237 milliseconds (best of 16 runs)

📝 Explanation and details

The optimized code achieves a 71% speedup (407ms → 237ms) by replacing TensorFlow's tensor_scatter_nd_update operations with a more efficient one-hot mask-based update strategy.

Key optimizations:

Eliminated expensive scatter operations: The original code called tf.tensor_scatter_nd_update twice per iteration, each requiring index tensor creation via tf.reshape(i, [1, 1]) and value reshaping via tf.reshape(c_val, [1]). The line profiler shows these scatter operations consumed ~23% of runtime (105ms + 32ms). The optimized version replaces this with vectorized arithmetic using tf.one_hot to create a mask, then updates via c_prime * inv_mask + mask * c_val. This mask-based approach is faster because it avoids the overhead of dynamic index construction and scatter's internal branching logic.
Explicit element access with tf.gather: Changed implicit indexing (e.g., c_prime[i - 1]) to explicit tf.gather calls. While this adds slight overhead for gather operations, it makes the computational graph more uniform and predictable for TensorFlow's optimizer, and works better with the mask-based update pattern.
Reduced graph complexity: By eliminating multiple reshape and scatter operations, the optimized code creates a simpler computation graph with fewer nodes. This reduces TensorFlow's internal dispatch overhead and memory allocation/deallocation cycles.

Performance characteristics from tests:

Speedup is consistent across all test cases (~54-108% faster)
Benefits scale well: small systems (size 2) see ~65% improvement, large systems (size 500) see ~57% improvement
Sequential iterations show particularly strong gains (108% faster in the 10-iteration test), suggesting the simpler graph structure compounds benefits when executed repeatedly
Works equally well across different numeric scenarios (large/small coefficients, mixed signs, edge cases)

Why this matters:
The function appears to be a body function for a loop-based tridiagonal solver. Since it's designed to be called iteratively (as evidenced by returning i + 1 and updated arrays), the per-iteration savings compound significantly. The mask-based update pattern is a well-known TensorFlow optimization that trades a small amount of redundant computation (updating all elements with a mask rather than just one) for much lower dispatching and memory management overhead—a favorable tradeoff in TensorFlow's execution model.

✅ 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%

🌀 Click to see Generated Regression Tests

import numpy as np
import tensorflow as tf

from code_to_optimize.sample_code import _tridiagonal_forward_body_tf

# ============================================================================
# BASIC TEST CASES - Test fundamental functionality under normal conditions
# ============================================================================


def test_basic_forward_step_single_iteration():
    """Test that the function correctly processes a single forward step."""
    # Setup: Create simple tridiagonal system with index 1 (second row)
    i = tf.constant(1, dtype=tf.int32)
    c_prime = tf.constant([0.0, 0.0, 0.0], dtype=tf.float32)
    d_prime = tf.constant([0.0, 0.0, 0.0], dtype=tf.float32)
    n = tf.constant(3, dtype=tf.int32)
    a = tf.constant([0.0, 1.0, 1.0], dtype=tf.float32)
    b = tf.constant([2.0, 2.0, 2.0], dtype=tf.float32)
    c = tf.constant([1.0, 1.0, 0.0], dtype=tf.float32)
    d = tf.constant([4.0, 4.0, 4.0], dtype=tf.float32)

    # Execute the function
    result_i, result_c_prime, result_d_prime, result_n, result_a, result_b, result_c, result_d = (
        _tridiagonal_forward_body_tf(i, c_prime, d_prime, n, a, b, c, d)
    )  # 15.0ms -> 9.09ms (65.3% faster)


def test_output_structure():
    """Test that the function returns correct number of outputs with correct types."""
    # Setup: Minimal valid inputs
    i = tf.constant(1, dtype=tf.int32)
    c_prime = tf.constant([0.0, 0.0], dtype=tf.float32)
    d_prime = tf.constant([0.0, 0.0], dtype=tf.float32)
    n = tf.constant(2, dtype=tf.int32)
    a = tf.constant([0.0, 1.0], dtype=tf.float32)
    b = tf.constant([2.0, 2.0], dtype=tf.float32)
    c = tf.constant([1.0, 0.0], dtype=tf.float32)
    d = tf.constant([4.0, 4.0], dtype=tf.float32)

    # Execute the function
    codeflash_output = _tridiagonal_forward_body_tf(i, c_prime, d_prime, n, a, b, c, d)
    result = codeflash_output  # 14.3ms -> 9.09ms (57.3% faster)


def test_basic_calculation_accuracy():
    """Test that c_prime and d_prime are calculated correctly."""
    # Setup: Simple case where calculation is easy to verify
    # Using: c_prime[i] = c[i] / (b[i] - a[i-1] * c_prime[i-1])
    #        d_prime[i] = (d[i] - a[i-1] * d_prime[i-1]) / (b[i] - a[i-1] * c_prime[i-1])
    i = tf.constant(1, dtype=tf.int32)
    c_prime = tf.constant([0.0, 0.0], dtype=tf.float64)
    d_prime = tf.constant([2.0, 0.0], dtype=tf.float64)  # d_prime[0] = 2.0
    n = tf.constant(2, dtype=tf.int32)
    a = tf.constant([0.0, 2.0], dtype=tf.float64)
    b = tf.constant([4.0, 4.0], dtype=tf.float64)
    c = tf.constant([1.0, 0.0], dtype=tf.float64)
    d = tf.constant([8.0, 8.0], dtype=tf.float64)

    # Execute
    _, result_c_prime, result_d_prime, _, _, _, _, _ = _tridiagonal_forward_body_tf(
        i, c_prime, d_prime, n, a, b, c, d
    )  # 15.0ms -> 9.10ms (64.9% faster)

    # Manual calculation:
    # denom = b[1] - a[0] * c_prime[0] = 4.0 - 0.0 * 0.0 = 4.0
    # c_prime[1] = c[1] / denom = 0.0 / 4.0 = 0.0
    # d_prime[1] = (d[1] - a[0] * d_prime[0]) / denom = (8.0 - 0.0 * 2.0) / 4.0 = 2.0
    expected_c_prime_1 = tf.constant(0.0, dtype=tf.float64)
    expected_d_prime_1 = tf.constant(2.0, dtype=tf.float64)


def test_float32_preservation():
    """Test that float32 dtype is preserved through computation."""
    # Setup: Create inputs with explicit float32 dtype
    i = tf.constant(1, dtype=tf.int32)
    c_prime = tf.constant([0.0, 0.0], dtype=tf.float32)
    d_prime = tf.constant([1.0, 0.0], dtype=tf.float32)
    n = tf.constant(2, dtype=tf.int32)
    a = tf.constant([0.0, 1.5], dtype=tf.float32)
    b = tf.constant([3.0, 3.0], dtype=tf.float32)
    c = tf.constant([0.5, 0.0], dtype=tf.float32)
    d = tf.constant([6.0, 6.0], dtype=tf.float32)

    # Execute
    _, result_c_prime, result_d_prime, _, _, _, _, _ = _tridiagonal_forward_body_tf(
        i, c_prime, d_prime, n, a, b, c, d
    )  # 15.0ms -> 9.09ms (64.7% faster)


def test_index_casting_from_python_int():
    """Test that the function casts Python int to int32 correctly."""
    # Setup: Pass Python int instead of tensor for i
    i = 1  # Python int
    c_prime = tf.constant([0.0, 0.0], dtype=tf.float32)
    d_prime = tf.constant([1.0, 0.0], dtype=tf.float32)
    n = tf.constant(2, dtype=tf.int32)
    a = tf.constant([0.0, 1.0], dtype=tf.float32)
    b = tf.constant([2.0, 2.0], dtype=tf.float32)
    c = tf.constant([1.0, 0.0], dtype=tf.float32)
    d = tf.constant([4.0, 4.0], dtype=tf.float32)

    # Execute - should not raise an error
    result_i, _, _, _, _, _, _, _ = _tridiagonal_forward_body_tf(
        i, c_prime, d_prime, n, a, b, c, d
    )  # 10.9ms -> 7.05ms (53.9% faster)


# ============================================================================
# EDGE TEST CASES - Test behavior under extreme or unusual conditions
# ============================================================================


def test_minimum_system_size():
    """Test with minimum valid system size (2 equations)."""
    # Setup: 2x2 tridiagonal system
    i = tf.constant(1, dtype=tf.int32)
    c_prime = tf.constant([0.0, 0.0], dtype=tf.float32)
    d_prime = tf.constant([1.0, 0.0], dtype=tf.float32)
    n = tf.constant(2, dtype=tf.int32)
    a = tf.constant([0.0, 1.0], dtype=tf.float32)
    b = tf.constant([2.0, 2.0], dtype=tf.float32)
    c = tf.constant([1.0, 0.0], dtype=tf.float32)
    d = tf.constant([4.0, 4.0], dtype=tf.float32)

    # Execute - should work without issues
    result_i, result_c_prime, result_d_prime, result_n, _, _, _, _ = _tridiagonal_forward_body_tf(
        i, c_prime, d_prime, n, a, b, c, d
    )  # 15.0ms -> 9.07ms (65.1% faster)


def test_zero_denominator_handling():
    """Test behavior when denominator approaches zero (near-singular system)."""
    # Setup: Create case where denom = b[i] - a[i-1] * c_prime[i-1] is very small
    # This tests numerical stability
    i = tf.constant(1, dtype=tf.int32)
    c_prime = tf.constant([0.5, 0.0], dtype=tf.float64)
    d_prime = tf.constant([1.0, 0.0], dtype=tf.float64)
    n = tf.constant(2, dtype=tf.int32)
    a = tf.constant([0.0, 1.0], dtype=tf.float64)
    b = tf.constant([0.5, 0.5], dtype=tf.float64)  # denom will be 0.5 - 1.0*0.5 = 0.0
    c = tf.constant([1.0, 0.0], dtype=tf.float64)
    d = tf.constant([4.0, 4.0], dtype=tf.float64)

    # Execute - TensorFlow may produce inf or nan
    _, result_c_prime, result_d_prime, _, _, _, _, _ = _tridiagonal_forward_body_tf(
        i, c_prime, d_prime, n, a, b, c, d
    )  # 15.0ms -> 9.06ms (65.1% faster)


def test_very_large_coefficient_values():
    """Test with very large coefficient values to check numerical stability."""
    # Setup: Large coefficient values
    i = tf.constant(1, dtype=tf.int32)
    c_prime = tf.constant([0.0, 0.0], dtype=tf.float64)
    d_prime = tf.constant([1e6, 0.0], dtype=tf.float64)
    n = tf.constant(2, dtype=tf.int32)
    a = tf.constant([0.0, 1e6], dtype=tf.float64)
    b = tf.constant([1e6, 1e6], dtype=tf.float64)
    c = tf.constant([1e6, 0.0], dtype=tf.float64)
    d = tf.constant([1e6, 1e6], dtype=tf.float64)

    # Execute
    _, result_c_prime, result_d_prime, _, _, _, _, _ = _tridiagonal_forward_body_tf(
        i, c_prime, d_prime, n, a, b, c, d
    )  # 15.0ms -> 9.05ms (65.8% faster)


def test_very_small_coefficient_values():
    """Test with very small coefficient values to check numerical stability."""
    # Setup: Small coefficient values
    i = tf.constant(1, dtype=tf.int32)
    c_prime = tf.constant([0.0, 0.0], dtype=tf.float64)
    d_prime = tf.constant([1e-6, 0.0], dtype=tf.float64)
    n = tf.constant(2, dtype=tf.int32)
    a = tf.constant([0.0, 1e-6], dtype=tf.float64)
    b = tf.constant([1e-6, 1e-6], dtype=tf.float64)
    c = tf.constant([1e-6, 0.0], dtype=tf.float64)
    d = tf.constant([1e-6, 1e-6], dtype=tf.float64)

    # Execute
    _, result_c_prime, result_d_prime, _, _, _, _, _ = _tridiagonal_forward_body_tf(
        i, c_prime, d_prime, n, a, b, c, d
    )  # 15.0ms -> 9.05ms (65.8% faster)


def test_negative_coefficient_values():
    """Test with negative coefficient values."""
    # Setup: Negative coefficients (valid in tridiagonal systems)
    i = tf.constant(1, dtype=tf.int32)
    c_prime = tf.constant([0.0, 0.0], dtype=tf.float32)
    d_prime = tf.constant([-1.0, 0.0], dtype=tf.float32)
    n = tf.constant(2, dtype=tf.int32)
    a = tf.constant([0.0, -1.0], dtype=tf.float32)
    b = tf.constant([-2.0, -2.0], dtype=tf.float32)
    c = tf.constant([-1.0, 0.0], dtype=tf.float32)
    d = tf.constant([-4.0, -4.0], dtype=tf.float32)

    # Execute
    _, result_c_prime, result_d_prime, _, _, _, _, _ = _tridiagonal_forward_body_tf(
        i, c_prime, d_prime, n, a, b, c, d
    )  # 15.0ms -> 8.99ms (67.1% faster)


def test_mixed_sign_coefficients():
    """Test with mixed positive and negative coefficient values."""
    # Setup: Mix of positive and negative values
    i = tf.constant(1, dtype=tf.int32)
    c_prime = tf.constant([0.5, 0.0], dtype=tf.float32)
    d_prime = tf.constant([2.0, 0.0], dtype=tf.float32)
    n = tf.constant(2, dtype=tf.int32)
    a = tf.constant([0.0, -1.5], dtype=tf.float32)
    b = tf.constant([2.0, 3.0], dtype=tf.float32)
    c = tf.constant([-0.5, 0.0], dtype=tf.float32)
    d = tf.constant([4.0, -2.0], dtype=tf.float32)

    # Execute
    _, result_c_prime, result_d_prime, _, _, _, _, _ = _tridiagonal_forward_body_tf(
        i, c_prime, d_prime, n, a, b, c, d
    )  # 15.0ms -> 9.06ms (65.5% faster)


def test_zero_c_prime_previous():
    """Test when previous c_prime value is exactly zero."""
    # Setup: c_prime[0] = 0
    i = tf.constant(1, dtype=tf.int32)
    c_prime = tf.constant([0.0, 0.0], dtype=tf.float32)
    d_prime = tf.constant([1.0, 0.0], dtype=tf.float32)
    n = tf.constant(2, dtype=tf.int32)
    a = tf.constant([0.0, 2.0], dtype=tf.float32)
    b = tf.constant([2.0, 2.0], dtype=tf.float32)
    c = tf.constant([1.0, 0.0], dtype=tf.float32)
    d = tf.constant([4.0, 4.0], dtype=tf.float32)

    # Execute
    _, result_c_prime, result_d_prime, _, _, _, _, _ = _tridiagonal_forward_body_tf(
        i, c_prime, d_prime, n, a, b, c, d
    )  # 15.0ms -> 9.09ms (65.3% faster)


def test_zero_a_coefficient():
    """Test when a[i-1] (lower diagonal) is zero."""
    # Setup: a[0] = 0 (no lower diagonal element)
    i = tf.constant(1, dtype=tf.int32)
    c_prime = tf.constant([0.5, 0.0], dtype=tf.float32)
    d_prime = tf.constant([1.0, 0.0], dtype=tf.float32)
    n = tf.constant(2, dtype=tf.int32)
    a = tf.constant([0.0, 0.0], dtype=tf.float32)
    b = tf.constant([2.0, 2.0], dtype=tf.float32)
    c = tf.constant([1.0, 0.0], dtype=tf.float32)
    d = tf.constant([4.0, 4.0], dtype=tf.float32)

    # Execute
    _, result_c_prime, result_d_prime, _, _, _, _, _ = _tridiagonal_forward_body_tf(
        i, c_prime, d_prime, n, a, b, c, d
    )  # 15.0ms -> 9.02ms (66.4% faster)


def test_zero_c_coefficient():
    """Test when c[i] (upper diagonal) is zero."""
    # Setup: c[1] = 0 (no upper diagonal element at last row)
    i = tf.constant(1, dtype=tf.int32)
    c_prime = tf.constant([0.5, 0.0], dtype=tf.float32)
    d_prime = tf.constant([1.0, 0.0], dtype=tf.float32)
    n = tf.constant(2, dtype=tf.int32)
    a = tf.constant([0.0, 1.0], dtype=tf.float32)
    b = tf.constant([2.0, 2.0], dtype=tf.float32)
    c = tf.constant([1.0, 0.0], dtype=tf.float32)
    d = tf.constant([4.0, 4.0], dtype=tf.float32)

    # Execute
    _, result_c_prime, result_d_prime, _, _, _, _, _ = _tridiagonal_forward_body_tf(
        i, c_prime, d_prime, n, a, b, c, d
    )  # 15.0ms -> 9.02ms (66.6% faster)


def test_int64_index_conversion():
    """Test that int64 index is properly converted to int32."""
    # Setup: Pass int64 tensor for i
    i = tf.constant(1, dtype=tf.int64)
    c_prime = tf.constant([0.0, 0.0], dtype=tf.float32)
    d_prime = tf.constant([1.0, 0.0], dtype=tf.float32)
    n = tf.constant(2, dtype=tf.int32)
    a = tf.constant([0.0, 1.0], dtype=tf.float32)
    b = tf.constant([2.0, 2.0], dtype=tf.float32)
    c = tf.constant([1.0, 0.0], dtype=tf.float32)
    d = tf.constant([4.0, 4.0], dtype=tf.float32)

    # Execute
    result_i, _, _, _, _, _, _, _ = _tridiagonal_forward_body_tf(
        i, c_prime, d_prime, n, a, b, c, d
    )  # 15.5ms -> 8.54ms (80.8% faster)


def test_float64_precision():
    """Test with float64 for higher precision calculations."""
    # Setup: Use float64 for better precision
    i = tf.constant(1, dtype=tf.int32)
    c_prime = tf.constant([0.0, 0.0], dtype=tf.float64)
    d_prime = tf.constant([1.0, 0.0], dtype=tf.float64)
    n = tf.constant(2, dtype=tf.int32)
    a = tf.constant([0.0, 1.0 / 3.0], dtype=tf.float64)
    b = tf.constant([3.0, 3.0], dtype=tf.float64)
    c = tf.constant([1.0 / 3.0, 0.0], dtype=tf.float64)
    d = tf.constant([4.0, 4.0], dtype=tf.float64)

    # Execute
    _, result_c_prime, result_d_prime, _, _, _, _, _ = _tridiagonal_forward_body_tf(
        i, c_prime, d_prime, n, a, b, c, d
    )  # 15.0ms -> 9.08ms (65.5% faster)


# ============================================================================
# LARGE SCALE TEST CASES - Test performance and scalability
# ============================================================================


def test_large_system_size_100():
    """Test with a larger tridiagonal system (size 100)."""
    # Setup: Create 100x100 tridiagonal system
    size = 100
    i = tf.constant(50, dtype=tf.int32)  # Test at middle index
    c_prime = tf.constant(np.zeros(size, dtype=np.float32), dtype=tf.float32)
    d_prime = tf.constant(np.ones(size, dtype=np.float32), dtype=tf.float32)
    n = tf.constant(size, dtype=tf.int32)
    a = tf.constant(np.full(size, 0.5, dtype=np.float32), dtype=tf.float32)
    b = tf.constant(np.full(size, 2.0, dtype=np.float32), dtype=tf.float32)
    c = tf.constant(np.full(size, 0.5, dtype=np.float32), dtype=tf.float32)
    d = tf.constant(np.full(size, 4.0, dtype=np.float32), dtype=tf.float32)

    # Execute
    result_i, result_c_prime, result_d_prime, result_n, _, _, _, _ = _tridiagonal_forward_body_tf(
        i, c_prime, d_prime, n, a, b, c, d
    )  # 14.3ms -> 9.09ms (57.3% faster)


def test_large_system_size_500():
    """Test with a very large tridiagonal system (size 500)."""
    # Setup: Create 500x500 tridiagonal system
    size = 500
    i = tf.constant(250, dtype=tf.int32)  # Test at middle index
    c_prime = tf.constant(np.zeros(size, dtype=np.float32), dtype=tf.float32)
    d_prime = tf.constant(np.random.randn(size).astype(np.float32), dtype=tf.float32)
    n = tf.constant(size, dtype=tf.int32)
    a = tf.constant(np.random.randn(size).astype(np.float32), dtype=tf.float32)
    b = tf.constant(np.random.randn(size).astype(np.float32) + 3.0, dtype=tf.float32)
    c = tf.constant(np.random.randn(size).astype(np.float32), dtype=tf.float32)
    d = tf.constant(np.random.randn(size).astype(np.float32), dtype=tf.float32)

    # Execute
    result_i, result_c_prime, result_d_prime, result_n, _, _, _, _ = _tridiagonal_forward_body_tf(
        i, c_prime, d_prime, n, a, b, c, d
    )  # 14.3ms -> 9.10ms (57.5% faster)


def test_large_system_sequential_iterations():
    """Test multiple sequential iterations on a large system."""
    # Setup: Create 100x100 system
    size = 100
    c_prime = tf.constant(np.zeros(size, dtype=np.float32), dtype=tf.float32)
    d_prime = tf.constant(np.ones(size, dtype=np.float32), dtype=tf.float32)
    n = tf.constant(size, dtype=tf.int32)
    a = tf.constant(np.full(size, 0.5, dtype=np.float32), dtype=tf.float32)
    b = tf.constant(np.full(size, 2.0, dtype=np.float32), dtype=tf.float32)
    c = tf.constant(np.full(size, 0.5, dtype=np.float32), dtype=tf.float32)
    d = tf.constant(np.full(size, 4.0, dtype=np.float32), dtype=tf.float32)

    # Execute multiple iterations
    current_i = tf.constant(1, dtype=tf.int32)
    for step in range(10):  # 10 iterations
        current_i, c_prime, d_prime, n, a, b, c, d = _tridiagonal_forward_body_tf(
            current_i, c_prime, d_prime, n, a, b, c, d
        )  # 63.7ms -> 30.7ms (108% faster)


def test_large_coefficients_large_system():
    """Test large system with large coefficient values."""
    # Setup: Large system with large coefficients
    size = 200
    i = tf.constant(100, dtype=tf.int32)
    c_prime = tf.constant(np.zeros(size, dtype=np.float64), dtype=tf.float64)
    d_prime = tf.constant(np.ones(size, dtype=np.float64) * 1e3, dtype=tf.float64)
    n = tf.constant(size, dtype=tf.int32)
    a = tf.constant(np.ones(size, dtype=np.float64) * 1e3, dtype=tf.float64)
    b = tf.constant(np.ones(size, dtype=np.float64) * 1e4, dtype=tf.float64)
    c = tf.constant(np.ones(size, dtype=np.float64) * 1e3, dtype=tf.float64)
    d = tf.constant(np.ones(size, dtype=np.float64) * 1e3, dtype=tf.float64)

    # Execute
    result_i, result_c_prime, result_d_prime, _, _, _, _, _ = _tridiagonal_forward_body_tf(
        i, c_prime, d_prime, n, a, b, c, d
    )  # 14.5ms -> 9.14ms (58.1% faster)


def test_array_modification_isolation():
    """Test that modifications don't leak between calls on large arrays."""
    # Setup: Create large array
    size = 150
    i1 = tf.constant(1, dtype=tf.int32)
    c_prime1 = tf.constant(np.zeros(size, dtype=np.float32), dtype=tf.float32)
    d_prime1 = tf.constant(np.ones(size, dtype=np.float32), dtype=tf.float32)
    n = tf.constant(size, dtype=tf.int32)
    a = tf.constant(np.full(size, 0.5, dtype=np.float32), dtype=tf.float32)
    b = tf.constant(np.full(size, 2.0, dtype=np.float32), dtype=tf.float32)
    c = tf.constant(np.full(size, 0.5, dtype=np.float32), dtype=tf.float32)
    d = tf.constant(np.full(size, 4.0, dtype=np.float32), dtype=tf.float32)

    # First call
    _, c_prime_after_1, d_prime_after_1, _, _, _, _, _ = _tridiagonal_forward_body_tf(
        i1, c_prime1, d_prime1, n, a, b, c, d
    )  # 14.2ms -> 9.10ms (56.5% faster)

    # Second call with same initial values should give same result
    i2 = tf.constant(1, dtype=tf.int32)
    c_prime2 = tf.constant(np.zeros(size, dtype=np.float32), dtype=tf.float32)
    d_prime2 = tf.constant(np.ones(size, dtype=np.float32), dtype=tf.float32)

    _, c_prime_after_2, d_prime_after_2, _, _, _, _, _ = _tridiagonal_forward_body_tf(
        i2, c_prime2, d_prime2, n, a, b, c, d
    )  # 5.48ms -> 2.33ms (135% faster)


def test_random_coefficients_large_system():
    """Test with random coefficients in a large system."""
    # Setup: Create large system with random coefficients
    size = 250
    np.random.seed(42)  # For reproducibility
    i = tf.constant(125, dtype=tf.int32)
    c_prime = tf.constant(np.random.randn(size).astype(np.float32), dtype=tf.float32)
    d_prime = tf.constant(np.random.randn(size).astype(np.float32), dtype=tf.float32)
    n = tf.constant(size, dtype=tf.int32)
    a = tf.constant(np.random.randn(size).astype(np.float32), dtype=tf.float32)
    # Ensure b has large magnitude to avoid singularity
    b = tf.constant((np.abs(np.random.randn(size)) + 5.0).astype(np.float32), dtype=tf.float32)
    c = tf.constant(np.random.randn(size).astype(np.float32), dtype=tf.float32)
    d = tf.constant(np.random.randn(size).astype(np.float32), dtype=tf.float32)

    # Execute
    result_i, result_c_prime, result_d_prime, _, _, _, _, _ = _tridiagonal_forward_body_tf(
        i, c_prime, d_prime, n, a, b, c, d
    )  # 14.3ms -> 9.08ms (57.5% faster)


def test_boundary_indices_large_system():
    """Test behavior at different boundary indices in large system."""
    # Setup: Create large system
    size = 300
    a = tf.constant(np.ones(size, dtype=np.float32) * 0.5, dtype=tf.float32)
    b = tf.constant(np.ones(size, dtype=np.float32) * 2.0, dtype=tf.float32)
    c = tf.constant(np.ones(size, dtype=np.float32) * 0.5, dtype=tf.float32)
    d = tf.constant(np.ones(size, dtype=np.float32) * 4.0, dtype=tf.float32)
    n = tf.constant(size, dtype=tf.int32)
    c_prime = tf.constant(np.zeros(size, dtype=np.float32), dtype=tf.float32)
    d_prime = tf.constant(np.ones(size, dtype=np.float32), dtype=tf.float32)

    # Test at different indices
    test_indices = [1, 50, 150, 299]
    for idx in test_indices:
        i = tf.constant(idx, dtype=tf.int32)
        result_i, result_c_prime, result_d_prime, _, _, _, _, _ = _tridiagonal_forward_body_tf(
            i, c_prime, d_prime, n, a, b, c, d
        )  # 30.6ms -> 16.0ms (91.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-_tridiagonal_forward_body_tf-mkgn5a91 and push.

The optimized code achieves a **71% speedup (407ms → 237ms)** by replacing TensorFlow's `tensor_scatter_nd_update` operations with a more efficient one-hot mask-based update strategy. **Key optimizations:** 1. **Eliminated expensive scatter operations**: The original code called `tf.tensor_scatter_nd_update` twice per iteration, each requiring index tensor creation via `tf.reshape(i, [1, 1])` and value reshaping via `tf.reshape(c_val, [1])`. The line profiler shows these scatter operations consumed ~23% of runtime (105ms + 32ms). The optimized version replaces this with vectorized arithmetic using `tf.one_hot` to create a mask, then updates via `c_prime * inv_mask + mask * c_val`. This mask-based approach is faster because it avoids the overhead of dynamic index construction and scatter's internal branching logic. 2. **Explicit element access with `tf.gather`**: Changed implicit indexing (e.g., `c_prime[i - 1]`) to explicit `tf.gather` calls. While this adds slight overhead for gather operations, it makes the computational graph more uniform and predictable for TensorFlow's optimizer, and works better with the mask-based update pattern. 3. **Reduced graph complexity**: By eliminating multiple reshape and scatter operations, the optimized code creates a simpler computation graph with fewer nodes. This reduces TensorFlow's internal dispatch overhead and memory allocation/deallocation cycles. **Performance characteristics from tests:** - Speedup is consistent across all test cases (~54-108% faster) - Benefits scale well: small systems (size 2) see ~65% improvement, large systems (size 500) see ~57% improvement - Sequential iterations show particularly strong gains (108% faster in the 10-iteration test), suggesting the simpler graph structure compounds benefits when executed repeatedly - Works equally well across different numeric scenarios (large/small coefficients, mixed signs, edge cases) **Why this matters:** The function appears to be a body function for a loop-based tridiagonal solver. Since it's designed to be called iteratively (as evidenced by returning `i + 1` and updated arrays), the per-iteration savings compound significantly. The mask-based update pattern is a well-known TensorFlow optimization that trades a small amount of redundant computation (updating all elements with a mask rather than just one) for much lower dispatching and memory management overhead—a favorable tradeoff in TensorFlow's execution model.

codeflash-ai · 2026-01-16T20:39:55Z

This PR has been automatically closed because the original PR #1055 by aseembits93 was closed.

codeflash-ai Bot requested a review from aseembits93 January 16, 2026 08:53

codeflash-ai Bot added ⚡️ codeflash Optimization PR opened by Codeflash AI 🎯 Quality: High Optimization Quality according to Codeflash labels Jan 16, 2026

Base automatically changed from instrument-jit to main January 16, 2026 20:39

codeflash-ai Bot closed this Jan 16, 2026

codeflash-ai Bot deleted the codeflash/optimize-_tridiagonal_forward_body_tf-mkgn5a91 branch January 16, 2026 20:39

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

⚡️ Speed up function `_tridiagonal_forward_body_tf` by 72%#1078

⚡️ Speed up function `_tridiagonal_forward_body_tf` by 72%#1078
codeflash-ai[bot] wants to merge 1 commit into
mainfrom
codeflash/optimize-_tridiagonal_forward_body_tf-mkgn5a91

codeflash-ai Bot commented Jan 16, 2026

Uh oh!

codeflash-ai Bot commented Jan 16, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

0 participants

Uh oh!

Conversation

codeflash-ai Bot commented Jan 16, 2026

📄 72% (0.72x) speedup for _tridiagonal_forward_body_tf in code_to_optimize/sample_code.py

📝 Explanation and details

Uh oh!

codeflash-ai Bot commented Jan 16, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

0 participants

📄 72% (0.72x) speedup for `_tridiagonal_forward_body_tf` in `code_to_optimize/sample_code.py`