⚡️ Speed up function _lis_inner_body_jax by 2,485%

[codeflash-ai]

📄 2,485% (24.85x) speedup for _lis_inner_body_jax in code_to_optimize/sample_code.py

⏱️ Runtime : 822 milliseconds → 31.8 milliseconds (best of 42 runs)

📝 Explanation and details

The optimized code achieves a 2484% speedup (from 822ms to 31.8ms) by adding JAX's @jit decorator to enable Just-In-Time compilation of the function.

Key Optimization

JIT Compilation via @jit decorator: The function performs array indexing, comparison operations, and conditional updates using JAX operations (jnp.where, .at[].set()). Without JIT, each of these operations is executed separately in Python with overhead for:

• Array indexing (arr[j], arr[i], dp_inner[j], dp_inner[i])
• Comparison operations (<, &, >)
• Conditional selection (jnp.where)
• Immutable array updates (dp_inner.at[i].set())

With @jit, JAX traces the function once and compiles it into optimized XLA code that:

1. Fuses operations: All operations are combined into a single compiled kernel, eliminating Python interpreter overhead between operations
2. Optimizes memory access patterns: Array accesses are optimized at the hardware level
3. Enables hardware acceleration: The compiled code can leverage GPU/TPU if available, or optimized CPU instructions

Why This Works

The function is a perfect candidate for JIT compilation because:

• It's a pure function with no side effects
• It uses only JAX array operations (not NumPy)
• It performs numerical computations that benefit from compiled execution
• The operation is relatively lightweight but called frequently (221 hits in profiler), making the compilation overhead worthwhile

Test Case Analysis

The speedup is consistent across all test scenarios:

• Simple updates: ~2000% speedup on basic operations
• Edge cases (equal values, negatives, zero values): ~1900-2300% speedup
• Large-scale tests: Even better gains (2586-2687%) when called in loops, as the JIT compilation cost is amortized over many calls

The optimization benefits any workload that calls this function repeatedly, particularly dynamic programming algorithms (like Longest Increasing Subsequence) where this inner body function would be invoked hundreds or thousands of times.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

import jax.numpy as jnp
import jax.random as jrandom

# imports
import pytest  # used for our unit tests

from code_to_optimize.sample_code import _lis_inner_body_jax

# unit tests


def test_basic_update_happens_when_condition_true():
    # Basic scenario: arr[j] < arr[i] and dp_inner[j] + 1 > dp_inner[i] -> update at i
    # Prepare inputs as jax arrays with integer dtype (consistent types for JIT)
    arr = jnp.array([1, 3], dtype=jnp.int32)  # arr[0] < arr[1]
    dp = jnp.array([1, 1], dtype=jnp.int32)  # dp[0] + 1 == 2 > dp[1] == 1
    # Call function: expect dp[1] to be updated to dp[0] + 1 == 2
    codeflash_output = _lis_inner_body_jax(0, dp, arr, 1)
    res = codeflash_output  # 3.68ms -> 178μs (1962% faster)


def test_no_update_when_arr_not_less():
    # Edge: arr[j] < arr[i] is false -> no update should occur
    arr = jnp.array([3, 1], dtype=jnp.int32)  # arr[0] < arr[1] is false
    dp = jnp.array([5, 2], dtype=jnp.int32)
    codeflash_output = _lis_inner_body_jax(0, dp, arr, 1)
    res = codeflash_output  # 3.62ms -> 169μs (2043% faster)


@pytest.mark.parametrize("dtype", [jnp.int32, jnp.int64])
def test_no_update_when_not_strictly_increasing(dtype):
    # Edge: arr[j] < arr[i] true, but dp[j] + 1 == dp[i] (not >) -> no update
    arr = jnp.array([1, 2], dtype=jnp.int32)  # values such that arr[0] < arr[1]
    dp = jnp.array([1, 2], dtype=dtype)  # dp[0] + 1 == 2 which is not > dp[1] == 2
    codeflash_output = _lis_inner_body_jax(0, dp, arr, 1)
    res = codeflash_output  # 7.15ms -> 340μs (2000% faster)


def test_equal_values_in_arr_no_update():
    # Edge: arr[j] == arr[i] -> arr[j] < arr[i] false -> no update
    arr = jnp.array([7, 7], dtype=jnp.int32)
    dp = jnp.array([0, 0], dtype=jnp.int32)
    codeflash_output = _lis_inner_body_jax(0, dp, arr, 1)
    res = codeflash_output  # 3.60ms -> 182μs (1879% faster)


def test_negative_dp_inner_values():
    # Edge: dp_inner contains negative values; arithmetic still holds
    # Choose arr so arr[0] < arr[1] true, and dp[0] + 1 > dp[1] (e.g., -3 + 1 = -2 > -10)
    arr = jnp.array([-5, -1], dtype=jnp.int32)
    dp = jnp.array([-3, -10], dtype=jnp.int32)
    codeflash_output = _lis_inner_body_jax(0, dp, arr, 1)
    res = codeflash_output  # 3.55ms -> 158μs (2136% faster)


def test_i_equals_j_no_update_and_safe_behavior():
    # Edge: i == j should be safe (arr[j] < arr[i] will be false since equal indices -> no update)
    arr = jnp.array([2, 4, 6], dtype=jnp.int32)
    dp = jnp.array([0, 1, 2], dtype=jnp.int32)
    # call with same index
    codeflash_output = _lis_inner_body_jax(1, dp, arr, 1)
    res = codeflash_output  # 3.68ms -> 174μs (2002% faster)


def test_float_dtype_behavior():
    # Verify the function works with floating-point dp_inner arrays as well
    arr = jnp.array([0.1, 0.2], dtype=jnp.float32)  # arr[0] < arr[1]
    dp = jnp.array([0.5, 0.0], dtype=jnp.float32)  # dp[0] + 1 = 1.5 > 0.0 -> update
    codeflash_output = _lis_inner_body_jax(0, dp, arr, 1)
    res = codeflash_output  # 3.65ms -> 163μs (2139% faster)


def test_large_scale_many_elements_but_limited_updates():
    # Large-scale style test: create a large array (but fewer than 1000 elements) and
    # perform a controlled number of updates (<<1000) to test scalability and behavior.
    size = 800  # large but under 1000 as requested
    # Create a strictly increasing array so arr[j] < arr[i] for j < i
    arr = jnp.arange(size, dtype=jnp.int32)
    # Initialize dp with zeros
    dp = jnp.zeros(size, dtype=jnp.int32)
    # Use a reproducible PRNG key from JAX
    key = jrandom.PRNGKey(1234)
    # We'll perform 100 random updates: choose random pairs (j < i)
    num_updates = 100
    # Split key for reproducibility
    key, subkey = jrandom.split(key)
    # Generate random i indices in [1, size-1]
    i_indices = jrandom.randint(subkey, shape=(num_updates,), minval=1, maxval=size, dtype=jnp.int32)
    key, subkey = jrandom.split(key)
    # For each i, choose j uniformly from [0, i-1] (we'll compute safe j using modulo to ensure j < i)
    raw_j = jrandom.randint(subkey, shape=(num_updates,), minval=0, maxval=size, dtype=jnp.int32)
    # Convert to Python lists for iteration (small loop <=1000 iterations)
    i_list = list(map(int, i_indices.tolist()))
    raw_j_list = list(map(int, raw_j.tolist()))
    dp_current = dp
    for idx in range(num_updates):
        i_val = i_list[idx]
        # ensure j < i by taking modulo i_val (i_val >= 1)
        j_val = raw_j_list[idx] % i_val
        # Call the function once per chosen pair
        codeflash_output = _lis_inner_body_jax(j_val, dp_current, arr, i_val)
        dp_current = codeflash_output  # 373ms -> 13.9ms (2586% faster)
    # After many updates on an increasing arr and starting from zeros, many dp entries should have become > 0
    dp_list = dp_current.tolist()
    positive_count = sum(1 for v in dp_list if v > 0)


def test_strictness_of_comparison_mutation_detection():
    # This test is intentionally precise to catch mutants that change '>' to '>=' or change arr comparison
    # Build values where dp[j] + 1 == dp[i] and arr[j] < arr[i] is True.
    arr = jnp.array([1, 2], dtype=jnp.int32)
    dp = jnp.array([1, 2], dtype=jnp.int32)  # dp[0] + 1 == dp[1]
    # No update expected because condition uses strict >; if the code used >= this test would fail.
    codeflash_output = _lis_inner_body_jax(0, dp, arr, 1)
    res = codeflash_output  # 3.70ms -> 179μs (1957% faster)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import jax.numpy as jnp
import numpy as np

# imports
from code_to_optimize.sample_code import _lis_inner_body_jax

# ============================================================================
# BASIC TEST CASES
# ============================================================================


def test_basic_update_when_condition_true():
    """Test that dp_inner[i] is updated when arr[j] < arr[i] and dp_inner[j] + 1 > dp_inner[i]"""
    arr = jnp.array([1.0, 3.0, 2.0], dtype=jnp.float32)
    dp_inner = jnp.array([1.0, 1.0, 1.0], dtype=jnp.float32)
    j = 0
    i = 2

    codeflash_output = _lis_inner_body_jax(j, dp_inner, arr, i)
    result = codeflash_output  # 3.76ms -> 174μs (2053% faster)

    # arr[0]=1.0 < arr[2]=2.0 is True, and dp_inner[0]+1=2.0 > dp_inner[2]=1.0 is True
    # So dp_inner[2] should be updated to 2.0
    expected = jnp.array([1.0, 1.0, 2.0], dtype=jnp.float32)


def test_basic_no_update_when_arr_condition_false():
    """Test that dp_inner[i] remains unchanged when arr[j] >= arr[i]"""
    arr = jnp.array([5.0, 3.0, 2.0], dtype=jnp.float32)
    dp_inner = jnp.array([1.0, 1.0, 1.0], dtype=jnp.float32)
    j = 0
    i = 2

    codeflash_output = _lis_inner_body_jax(j, dp_inner, arr, i)
    result = codeflash_output  # 3.68ms -> 158μs (2217% faster)

    # arr[0]=5.0 < arr[2]=2.0 is False, so condition is False
    # dp_inner[2] should remain 1.0
    expected = jnp.array([1.0, 1.0, 1.0], dtype=jnp.float32)


def test_basic_no_update_when_dp_condition_false():
    """Test that dp_inner[i] remains unchanged when dp_inner[j] + 1 <= dp_inner[i]"""
    arr = jnp.array([1.0, 3.0, 2.0], dtype=jnp.float32)
    dp_inner = jnp.array([1.0, 5.0, 10.0], dtype=jnp.float32)
    j = 0
    i = 2

    codeflash_output = _lis_inner_body_jax(j, dp_inner, arr, i)
    result = codeflash_output  # 3.72ms -> 161μs (2198% faster)

    # arr[0]=1.0 < arr[2]=2.0 is True, but dp_inner[0]+1=2.0 > dp_inner[2]=10.0 is False
    # So condition is False, dp_inner[2] should remain 10.0
    expected = jnp.array([1.0, 5.0, 10.0], dtype=jnp.float32)


def test_basic_single_element_update():
    """Test with minimal arrays of size 2"""
    arr = jnp.array([1.0, 2.0], dtype=jnp.float32)
    dp_inner = jnp.array([1.0, 1.0], dtype=jnp.float32)
    j = 0
    i = 1

    codeflash_output = _lis_inner_body_jax(j, dp_inner, arr, i)
    result = codeflash_output  # 3.67ms -> 169μs (2069% faster)

    # arr[0]=1.0 < arr[1]=2.0 is True, dp_inner[0]+1=2.0 > dp_inner[1]=1.0 is True
    expected = jnp.array([1.0, 2.0], dtype=jnp.float32)


def test_basic_preserves_other_elements():
    """Test that only dp_inner[i] is modified, not other elements"""
    arr = jnp.array([1.0, 2.0, 3.0, 4.0], dtype=jnp.float32)
    dp_inner = jnp.array([1.0, 2.0, 3.0, 4.0], dtype=jnp.float32)
    j = 1
    i = 3

    codeflash_output = _lis_inner_body_jax(j, dp_inner, arr, i)
    result = codeflash_output  # 3.72ms -> 165μs (2141% faster)


# ============================================================================
# EDGE CASES
# ============================================================================


def test_edge_equal_values_in_arr():
    """Test when arr[j] == arr[i] (boundary condition for arr[j] < arr[i])"""
    arr = jnp.array([2.0, 2.0, 3.0], dtype=jnp.float32)
    dp_inner = jnp.array([1.0, 1.0, 1.0], dtype=jnp.float32)
    j = 0
    i = 1

    codeflash_output = _lis_inner_body_jax(j, dp_inner, arr, i)
    result = codeflash_output  # 3.72ms -> 165μs (2150% faster)

    # arr[0]=2.0 < arr[1]=2.0 is False, so no update
    expected = jnp.array([1.0, 1.0, 1.0], dtype=jnp.float32)


def test_edge_equal_dp_values():
    """Test when dp_inner[j] + 1 == dp_inner[i] (boundary condition for dp_inner[j] + 1 > dp_inner[i])"""
    arr = jnp.array([1.0, 3.0, 2.0], dtype=jnp.float32)
    dp_inner = jnp.array([1.0, 1.0, 2.0], dtype=jnp.float32)
    j = 0
    i = 2

    codeflash_output = _lis_inner_body_jax(j, dp_inner, arr, i)
    result = codeflash_output  # 3.71ms -> 172μs (2045% faster)

    # arr[0]=1.0 < arr[2]=2.0 is True, but dp_inner[0]+1=2.0 > dp_inner[2]=2.0 is False
    expected = jnp.array([1.0, 1.0, 2.0], dtype=jnp.float32)


def test_edge_zero_values_in_arr():
    """Test with zero values in array"""
    arr = jnp.array([0.0, 1.0, 0.5], dtype=jnp.float32)
    dp_inner = jnp.array([1.0, 1.0, 1.0], dtype=jnp.float32)
    j = 0
    i = 2

    codeflash_output = _lis_inner_body_jax(j, dp_inner, arr, i)
    result = codeflash_output  # 3.65ms -> 179μs (1936% faster)

    # arr[0]=0.0 < arr[2]=0.5 is True, dp_inner[0]+1=2.0 > dp_inner[2]=1.0 is True
    expected = jnp.array([1.0, 1.0, 2.0], dtype=jnp.float32)


def test_edge_negative_values_in_arr():
    """Test with negative values in array"""
    arr = jnp.array([-5.0, -2.0, -3.0], dtype=jnp.float32)
    dp_inner = jnp.array([1.0, 1.0, 1.0], dtype=jnp.float32)
    j = 0
    i = 2

    codeflash_output = _lis_inner_body_jax(j, dp_inner, arr, i)
    result = codeflash_output  # 3.66ms -> 157μs (2230% faster)

    # arr[0]=-5.0 < arr[2]=-3.0 is True, dp_inner[0]+1=2.0 > dp_inner[2]=1.0 is True
    expected = jnp.array([1.0, 1.0, 2.0], dtype=jnp.float32)


def test_edge_mixed_positive_negative():
    """Test with mix of positive and negative values"""
    arr = jnp.array([-1.0, 0.0, 1.0, -2.0], dtype=jnp.float32)
    dp_inner = jnp.array([1.0, 1.0, 1.0, 1.0], dtype=jnp.float32)
    j = 0
    i = 2

    codeflash_output = _lis_inner_body_jax(j, dp_inner, arr, i)
    result = codeflash_output  # 3.71ms -> 164μs (2157% faster)

    # arr[0]=-1.0 < arr[2]=1.0 is True, dp_inner[0]+1=2.0 > dp_inner[2]=1.0 is True
    expected = jnp.array([1.0, 1.0, 2.0, 1.0], dtype=jnp.float32)


def test_edge_j_equals_i():
    """Test behavior when j == i (comparing element with itself)"""
    arr = jnp.array([1.0, 2.0, 3.0], dtype=jnp.float32)
    dp_inner = jnp.array([1.0, 1.0, 1.0], dtype=jnp.float32)
    j = 1
    i = 1

    codeflash_output = _lis_inner_body_jax(j, dp_inner, arr, i)
    result = codeflash_output  # 3.68ms -> 160μs (2193% faster)

    # arr[1]=2.0 < arr[1]=2.0 is False, so no update
    expected = jnp.array([1.0, 1.0, 1.0], dtype=jnp.float32)


def test_edge_large_dp_values():
    """Test with very large dp_inner values"""
    arr = jnp.array([1.0, 2.0], dtype=jnp.float32)
    dp_inner = jnp.array([1e6, 1.0], dtype=jnp.float32)
    j = 0
    i = 1

    codeflash_output = _lis_inner_body_jax(j, dp_inner, arr, i)
    result = codeflash_output  # 3.65ms -> 161μs (2151% faster)

    # arr[0]=1.0 < arr[1]=2.0 is True, dp_inner[0]+1 > dp_inner[1] is True
    expected = jnp.array([1e6, 1e6 + 1], dtype=jnp.float32)


def test_edge_very_small_dp_values():
    """Test with very small positive dp_inner values"""
    arr = jnp.array([1.0, 2.0], dtype=jnp.float32)
    dp_inner = jnp.array([1e-6, 1e-7], dtype=jnp.float32)
    j = 0
    i = 1

    codeflash_output = _lis_inner_body_jax(j, dp_inner, arr, i)
    result = codeflash_output  # 3.65ms -> 159μs (2191% faster)

    # arr[0]=1.0 < arr[1]=2.0 is True, dp_inner[0]+1 > dp_inner[1] is True
    expected = jnp.array([1e-6, 1e-6 + 1], dtype=jnp.float32)


def test_edge_float64_precision():
    """Test with float64 dtype for higher precision"""
    arr = jnp.array([1.0, 2.0, 3.0], dtype=jnp.float64)
    dp_inner = jnp.array([1.0, 1.0, 1.0], dtype=jnp.float64)
    j = 0
    i = 2

    codeflash_output = _lis_inner_body_jax(j, dp_inner, arr, i)
    result = codeflash_output  # 3.68ms -> 155μs (2266% faster)

    expected = jnp.array([1.0, 1.0, 2.0], dtype=jnp.float64)


def test_edge_index_at_boundaries():
    """Test with indices at array boundaries"""
    arr = jnp.array([1.0, 2.0, 3.0, 4.0, 5.0], dtype=jnp.float32)
    dp_inner = jnp.array([1.0, 1.0, 1.0, 1.0, 1.0], dtype=jnp.float32)
    j = 0
    i = 4

    codeflash_output = _lis_inner_body_jax(j, dp_inner, arr, i)
    result = codeflash_output  # 3.67ms -> 175μs (1989% faster)

    # arr[0]=1.0 < arr[4]=5.0 is True, dp_inner[0]+1=2.0 > dp_inner[4]=1.0 is True
    expected = jnp.array([1.0, 1.0, 1.0, 1.0, 2.0], dtype=jnp.float32)


def test_edge_all_same_values_arr():
    """Test when all array values are identical"""
    arr = jnp.array([5.0, 5.0, 5.0, 5.0], dtype=jnp.float32)
    dp_inner = jnp.array([1.0, 2.0, 3.0, 4.0], dtype=jnp.float32)
    j = 0
    i = 3

    codeflash_output = _lis_inner_body_jax(j, dp_inner, arr, i)
    result = codeflash_output  # 3.69ms -> 170μs (2067% faster)

    # arr[0]=5.0 < arr[3]=5.0 is False, so no update
    expected = jnp.array([1.0, 2.0, 3.0, 4.0], dtype=jnp.float32)


def test_edge_all_same_values_dp():
    """Test when all dp_inner values are identical"""
    arr = jnp.array([1.0, 2.0, 3.0, 4.0], dtype=jnp.float32)
    dp_inner = jnp.array([5.0, 5.0, 5.0, 5.0], dtype=jnp.float32)
    j = 0
    i = 3

    codeflash_output = _lis_inner_body_jax(j, dp_inner, arr, i)
    result = codeflash_output  # 3.68ms -> 164μs (2137% faster)

    # arr[0]=1.0 < arr[3]=4.0 is True, dp_inner[0]+1=6.0 > dp_inner[3]=5.0 is True
    expected = jnp.array([5.0, 5.0, 5.0, 6.0], dtype=jnp.float32)


def test_edge_strictly_decreasing_arr():
    """Test with strictly decreasing array values"""
    arr = jnp.array([10.0, 8.0, 6.0, 4.0, 2.0], dtype=jnp.float32)
    dp_inner = jnp.array([1.0, 1.0, 1.0, 1.0, 1.0], dtype=jnp.float32)
    j = 0
    i = 4

    codeflash_output = _lis_inner_body_jax(j, dp_inner, arr, i)
    result = codeflash_output  # 3.71ms -> 165μs (2147% faster)

    # arr[0]=10.0 < arr[4]=2.0 is False, so no update
    expected = jnp.array([1.0, 1.0, 1.0, 1.0, 1.0], dtype=jnp.float32)


def test_edge_strictly_increasing_arr():
    """Test with strictly increasing array values"""
    arr = jnp.array([1.0, 2.0, 3.0, 4.0, 5.0], dtype=jnp.float32)
    dp_inner = jnp.array([1.0, 1.0, 1.0, 1.0, 1.0], dtype=jnp.float32)
    j = 0
    i = 4

    codeflash_output = _lis_inner_body_jax(j, dp_inner, arr, i)
    result = codeflash_output  # 3.68ms -> 162μs (2166% faster)

    # arr[0]=1.0 < arr[4]=5.0 is True, dp_inner[0]+1=2.0 > dp_inner[4]=1.0 is True
    expected = jnp.array([1.0, 1.0, 1.0, 1.0, 2.0], dtype=jnp.float32)


def test_edge_extreme_negative_values():
    """Test with extreme negative values"""
    arr = jnp.array([-1e6, -1e5, -1e4], dtype=jnp.float32)
    dp_inner = jnp.array([1.0, 1.0, 1.0], dtype=jnp.float32)
    j = 0
    i = 2

    codeflash_output = _lis_inner_body_jax(j, dp_inner, arr, i)
    result = codeflash_output  # 3.69ms -> 161μs (2177% faster)

    # arr[0]=-1e6 < arr[2]=-1e4 is True, dp_inner[0]+1=2.0 > dp_inner[2]=1.0 is True
    expected = jnp.array([1.0, 1.0, 2.0], dtype=jnp.float32)


def test_edge_extreme_positive_values():
    """Test with extreme positive values"""
    arr = jnp.array([1e4, 1e5, 1e6], dtype=jnp.float32)
    dp_inner = jnp.array([1.0, 1.0, 1.0], dtype=jnp.float32)
    j = 0
    i = 2

    codeflash_output = _lis_inner_body_jax(j, dp_inner, arr, i)
    result = codeflash_output  # 3.70ms -> 159μs (2225% faster)

    # arr[0]=1e4 < arr[2]=1e6 is True, dp_inner[0]+1=2.0 > dp_inner[2]=1.0 is True
    expected = jnp.array([1.0, 1.0, 2.0], dtype=jnp.float32)


# ============================================================================
# LARGE SCALE TEST CASES
# ============================================================================


def test_large_scale_medium_array():
    """Test with moderately large array (size 100)"""
    arr_size = 100
    arr = jnp.arange(1.0, arr_size + 1.0, dtype=jnp.float32)
    dp_inner = jnp.ones(arr_size, dtype=jnp.float32)
    j = 0
    i = 99

    codeflash_output = _lis_inner_body_jax(j, dp_inner, arr, i)
    result = codeflash_output  # 3.64ms -> 169μs (2047% faster)

    # arr[0]=1.0 < arr[99]=100.0 is True, dp_inner[0]+1=2.0 > dp_inner[99]=1.0 is True
    expected = dp_inner.at[99].set(2.0)


def test_large_scale_various_indices():
    """Test with various index combinations in a large array"""
    arr_size = 200
    arr = jnp.arange(1.0, arr_size + 1.0, dtype=jnp.float32)
    dp_inner = jnp.ones(arr_size, dtype=jnp.float32)

    # Test multiple index pairs
    test_pairs = [(0, 50), (10, 100), (50, 199), (100, 150)]

    for j, i in test_pairs:
        codeflash_output = _lis_inner_body_jax(j, dp_inner, arr, i)
        result = codeflash_output  # 14.7ms -> 668μs (2103% faster)
        # arr[j] < arr[i] is always True for j < i in increasing sequence
        # dp_inner[j]+1 > dp_inner[i] is always True when all start at 1.0
        expected = dp_inner.at[i].set(2.0)


def test_large_scale_random_array_values():
    """Test with large array of random values"""
    arr_size = 256
    np.random.seed(42)
    arr = jnp.array(np.random.rand(arr_size).astype(np.float32))
    dp_inner = jnp.ones(arr_size, dtype=jnp.float32)
    j = 50
    i = 150

    codeflash_output = _lis_inner_body_jax(j, dp_inner, arr, i)
    result = codeflash_output  # 3.75ms -> 170μs (2100% faster)

    # Check that only index i is potentially modified
    for idx in range(arr_size):
        if idx != i:
            pass


def test_large_scale_dense_operations():
    """Test multiple operations on same large array to simulate inner loop"""
    arr_size = 128
    arr = jnp.arange(1.0, arr_size + 1.0, dtype=jnp.float32)
    dp_inner = jnp.ones(arr_size, dtype=jnp.float32)

    # Simulate multiple iterations
    result = dp_inner
    for j in range(10):
        i = arr_size - 1
        codeflash_output = _lis_inner_body_jax(j, result, arr, i)
        result = codeflash_output  # 37.0ms -> 1.49ms (2381% faster)


def test_large_scale_preserves_array_structure():
    """Test that large arrays maintain correct structure and dtype"""
    arr_size = 512
    arr = jnp.arange(1.0, arr_size + 1.0, dtype=jnp.float32)
    dp_inner = jnp.ones(arr_size, dtype=jnp.float32)
    j = 100
    i = 300

    codeflash_output = _lis_inner_body_jax(j, dp_inner, arr, i)
    result = codeflash_output  # 3.75ms -> 166μs (2159% faster)


def test_large_scale_performance_with_consecutive_updates():
    """Test performance with consecutive index updates"""
    arr_size = 256
    arr = jnp.arange(1.0, arr_size + 1.0, dtype=jnp.float32)
    dp_inner = jnp.ones(arr_size, dtype=jnp.float32)

    # Simulate consecutive updates from same j
    result = dp_inner
    j = 0
    for i in range(1, min(50, arr_size)):
        codeflash_output = _lis_inner_body_jax(j, result, arr, i)
        result = codeflash_output  # 186ms -> 6.70ms (2687% faster)

    # All indices from 1 to 49 should be updated to 2.0
    for i in range(1, 50):
        pass


def test_large_scale_high_dp_values():
    """Test with large dp_inner values"""
    arr_size = 100
    arr = jnp.arange(1.0, arr_size + 1.0, dtype=jnp.float32)
    dp_inner = jnp.full(arr_size, 500.0, dtype=jnp.float32)
    j = 0
    i = 99

    codeflash_output = _lis_inner_body_jax(j, dp_inner, arr, i)
    result = codeflash_output  # 3.66ms -> 161μs (2170% faster)

    # dp_inner[0]+1=501.0 > dp_inner[99]=500.0 is True
    expected = dp_inner.at[99].set(501.0)


def test_large_scale_alternating_pattern():
    """Test with alternating increase/decrease pattern in large array"""
    arr_size = 200
    # Create pattern: 1, 100, 2, 99, 3, 98, ...
    arr_vals = []
    for i in range(arr_size // 2):
        arr_vals.append(float(i + 1))
        arr_vals.append(float(arr_size - i))
    arr = jnp.array(arr_vals[:arr_size], dtype=jnp.float32)
    dp_inner = jnp.ones(arr_size, dtype=jnp.float32)

    j = 0
    i = arr_size - 1
    codeflash_output = _lis_inner_body_jax(j, dp_inner, arr, i)
    result = codeflash_output  # 3.55ms -> 158μs (2135% faster)


def test_large_scale_backwards_iteration():
    """Test updating from near-end indices"""
    arr_size = 150
    arr = jnp.arange(1.0, arr_size + 1.0, dtype=jnp.float32)
    dp_inner = jnp.ones(arr_size, dtype=jnp.float32)

    # Update from high j to lower i (which would be unusual but should still work)
    j = arr_size - 2
    i = arr_size - 1

    codeflash_output = _lis_inner_body_jax(j, dp_inner, arr, i)
    result = codeflash_output  # 3.66ms -> 167μs (2087% faster)

    # arr[j] < arr[i] is True, dp_inner[j]+1 > dp_inner[i] is True
    expected = dp_inner.at[i].set(2.0)


def test_large_scale_uniform_large_values():
    """Test with large uniform array values"""
    arr_size = 100
    arr = jnp.full(arr_size, 1e5, dtype=jnp.float32)
    # Modify one element to make comparison work
    arr = arr.at[0].set(1e5 - 1)
    dp_inner = jnp.ones(arr_size, dtype=jnp.float32)
    j = 0
    i = 99

    codeflash_output = _lis_inner_body_jax(j, dp_inner, arr, i)
    result = codeflash_output  # 3.56ms -> 148μs (2306% faster)

    # arr[0] < arr[99] is True, dp_inner[0]+1 > dp_inner[99] is True
    expected = dp_inner.at[99].set(2.0)


def test_large_scale_immutability_verification():
    """Verify that the original arrays are not modified"""
    arr_size = 200
    arr = jnp.arange(1.0, arr_size + 1.0, dtype=jnp.float32)
    dp_inner = jnp.ones(arr_size, dtype=jnp.float32)
    arr_original = jnp.array(arr)
    dp_original = jnp.array(dp_inner)

    j = 50
    i = 150
    codeflash_output = _lis_inner_body_jax(j, dp_inner, arr, i)
    result = codeflash_output  # 3.51ms -> 150μs (2227% faster)


def test_large_scale_stability_multiple_updates():
    """Test numerical stability across multiple update operations"""
    arr_size = 100
    arr = jnp.arange(1.0, arr_size + 1.0, dtype=jnp.float32)
    dp_inner = jnp.ones(arr_size, dtype=jnp.float32)

    # Apply multiple updates sequentially
    result = dp_inner
    for _ in range(20):
        codeflash_output = _lis_inner_body_jax(0, result, arr, 99)
        result = codeflash_output  # 71.0ms -> 2.74ms (2490% 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-_lis_inner_body_jax-mkgiirxu and push.