⚡️ Speed up function demo_postprocess by 39%

[codeflash-ai]

📄 39% (0.39x) speedup for demo_postprocess in unstructured_inference/models/yolox.py

⏱️ Runtime : 76.3 milliseconds → 54.8 milliseconds (best of 5 runs)

📝 Explanation and details

The optimized code achieves a 39% speedup through two key optimizations:

1. Grid Caching (~9.5% of runtime)

The original code rebuilt meshgrids from scratch on every call via np.meshgrid(np.arange(wsize), np.arange(hsize)), taking ~15.3% of runtime. The optimization introduces _GRID_CACHE to memoize these grids by (hsize, wsize) pairs. Since demo_postprocess is called repeatedly with the same input shape (1024, 768) in the inference pipeline (see function_references), subsequent calls retrieve cached grids in ~0.1-0.2μs vs ~93μs to rebuild, eliminating redundant computation.

2. Eliminating Intermediate Array Concatenation (~88% of runtime)

The original code built full grids and expanded_strides arrays via:

• Multiple np.concatenate() calls (1.8% runtime)
• Broadcasting these large arrays across the entire outputs tensor in vectorized operations (74.4% runtime for the two broadcast multiplications)

The optimization replaces this with per-stride slice processing: it iterates through each stride block, directly updating the corresponding slice of outputs in-place. This avoids:

• Allocating temporary arrays (grids: 1×8400×2, expanded_strides: 1×8400×1 for 1024×768 images)
• Broadcasting these arrays across the full tensor
• Memory copies during concatenation

Instead, each stride's computation uses only a small cached grid (e.g., 1×16384×2 for stride=8) that's added/multiplied with the relevant slice. This reduces peak memory usage and cache thrashing, particularly beneficial for large image sizes (test results show 42.9% speedup for 1280×1280 images).

Performance Characteristics

• Best for: Repeated calls with identical img_size (caching maximizes benefit) and large images/batches (avoids expensive concatenation overhead). Test cases show 81-329% speedups for such scenarios.
• Workload Impact: Since demo_postprocess is called in the hot path of YOLOX inference (once per image in image_processing), this 39% reduction directly improves end-to-end detection throughput. The optimization is especially impactful for batch processing or video streams where the same input shape recurs.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

import math  # used for numerical checks with math.isclose and math.exp

import numba  # required by the function decorator (must be present for the original function)
import numpy as np  # used to construct array inputs
import pytest  # used for our unit tests
from unstructured_inference.models.yolox import demo_postprocess

# unit tests

def _compute_grid_and_strides(img_size, p6=False):
    """
    Helper to compute hsizes, wsizes, total, grids and expanded_strides
    using the same logic as demo_postprocess (but implemented in pure python).
    This is used to build expected values for assertions.
    """
    if p6:
        strides = (8, 16, 32, 64)
    else:
        strides = (8, 16, 32)

    hsizes = [img_size[0] // s for s in strides]
    wsizes = [img_size[1] // s for s in strides]
    total = sum(h * w for h, w in zip(hsizes, wsizes))

    grids = np.empty((1, total, 2), dtype=np.float64)
    expanded_strides = np.empty((1, total, 1), dtype=np.float64)
    offset = 0
    for idx, stride in enumerate(strides):
        h = hsizes[idx]
        w = wsizes[idx]
        for y in range(h):
            for x in range(w):
                pos = offset + y * w + x
                grids[0, pos, 0] = x
                grids[0, pos, 1] = y
                expanded_strides[0, pos, 0] = stride
        offset += h * w

    return hsizes, wsizes, total, grids, expanded_strides

def _close(a, b, rel_tol=1e-7, abs_tol=1e-9):
    """Tiny wrapper around math.isclose for readability."""
    return math.isclose(a, b, rel_tol=rel_tol, abs_tol=abs_tol)

def test_basic_functionality_single_batch_zero_inputs():
    # Basic case: single batch, small image, zero-initialized outputs.
    img_size = (32, 32)  # small size to keep total small
    hsizes, wsizes, total, grids, expanded_strides = _compute_grid_and_strides(img_size, p6=False)

    # create outputs filled with zeros; channels = 4 (x, y, w, h)
    outputs = np.zeros((1, total, 4), dtype=np.float64)

    # call the function; it mutates `outputs` in-place and returns it
    codeflash_output = demo_postprocess(outputs, img_size, p6=False); result = codeflash_output # 194μs -> 68.7μs (183% faster)

    # Check first cell (pos 0) corresponds to first stride (8) and grid (0,0)
    stride0 = int(expanded_strides[0, 0, 0])

    # Check last cell of the first stride block (idx 0 block has hsizes[0]*wsizes[0] entries)
    first_block_count = int(hsizes[0] * wsizes[0])
    last_pos_first_block = first_block_count - 1
    # compute expected grid coords for that position: x = w-1, y = h-1 for that block
    w0 = int(wsizes[0])
    h0 = int(hsizes[0])
    expected_x = (w0 - 1) * stride0
    expected_y = (h0 - 1) * stride0

def test_multi_batch_leading_dims_different_initial_values():
    # Edge case: multiple batches (outer > 1) are supported; ensure independent processing.
    img_size = (32, 32)
    hsizes, wsizes, total, grids, expanded_strides = _compute_grid_and_strides(img_size, p6=False)

    # Build outputs with two "batches" (outer=2). Use channels >=4.
    # Use different initial x values in each batch to ensure outputs differ accordingly after processing.
    outputs = np.zeros((2, total, 4), dtype=np.float64)
    # Set x channel of second batch to 1.0 so we can detect the difference after postprocess
    outputs[1, :, 0] = 1.0

    # Copy for later manual computation
    original = outputs.copy()

    # Call the function
    demo_postprocess(outputs, img_size, p6=False) # 187μs -> 56.8μs (230% faster)

    # For position 0 (first grid cell), gridx == 0 and stride is expanded_strides[0,0,0]
    stride0 = int(expanded_strides[0, 0, 0])

    # The difference between batches at pos 0 should exactly equal stride0
    diff = outputs[1, 0, 0] - outputs[0, 0, 0]

def test_p6_true_adds_extra_stride_block():
    # Edge case: p6=True should include an additional stride (64)
    img_size = (64, 64)
    hsizes, wsizes, total, grids, expanded_strides = _compute_grid_and_strides(img_size, p6=True)

    # outputs zeros so w,h become equal to respective stride values after processing
    outputs = np.zeros((1, total, 4), dtype=np.float64)

    # Run the function with p6 True; ensures the code path that sets 4 strides executes.
    demo_postprocess(outputs, img_size, p6=True) # 235μs -> 78.2μs (201% faster)

    # The last block corresponds to the largest stride (64). Find a position within that block.
    # Compute where the last block starts:
    # For p6, strides are (8,16,32,64). The last block size is hsizes[-1] * wsizes[-1]
    last_block_count = int(hsizes[-1] * wsizes[-1])

    # pick the final position (last cell in the entire list)
    last_pos = total - 1
    last_stride = int(expanded_strides[0, last_pos, 0])

def test_large_scale_random_subset_verification():
    # Large-scale test (but within limits): use img_size which yields total < 1000 to satisfy constraints.
    img_size = (128, 128)  # yields totals that are moderate (kept under 1000)
    hsizes, wsizes, total, grids, expanded_strides = _compute_grid_and_strides(img_size, p6=False)

    # Use a deterministic RNG for reproducibility
    rng = np.random.RandomState(12345)

    # create outputs with random values for channels >= 4 (we'll include some extra channels)
    channels = 6
    outputs = rng.randn(1, total, channels).astype(np.float64)

    # Keep a copy of the original values for manual expected computation
    original = outputs.copy()

    # Run the function (mutates `outputs`)
    demo_postprocess(outputs, img_size, p6=False) # 222μs -> 68.2μs (226% faster)

    # Check a representative subset (first 20 cells) to ensure correct transformation
    check_count = min(20, total)
    for pos in range(check_count):
        # compute stride for this position
        stride = float(expanded_strides[0, pos, 0])
        grid_x = float(grids[0, pos, 0])
        grid_y = float(grids[0, pos, 1])

        # original x,y at [0,pos,0] and [0,pos,1]
        orig_x = float(original[0, pos, 0])
        orig_y = float(original[0, pos, 1])
        # after processing expected:
        expected_x = (orig_x + grid_x) * stride
        expected_y = (orig_y + grid_y) * stride

        # check width and height transformed by exponential then stride
        orig_w = float(original[0, pos, 2])
        orig_h = float(original[0, pos, 3])
        expected_w = math.exp(orig_w) * stride
        expected_h = math.exp(orig_h) * stride
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import numba
import numpy as np
import pytest
from unstructured_inference.models.yolox import demo_postprocess

class TestDemoPostprocessBasic:
    """Basic test cases for demo_postprocess function with normal conditions."""
    
    def test_basic_single_stride_no_p6(self):
        """Test basic functionality with standard p6=False configuration (8, 16, 32 strides)."""
        # Create a simple output tensor with shape (1, 1215, 85)
        # 1215 = (40*40 + 20*20 + 10*10) grid cells for 320x320 image
        img_size = (320, 320)
        batch_size = 1
        num_classes = 80
        channels = 5 + num_classes  # 4 coords + 1 objectness + num_classes
        
        # Calculate expected total anchors
        total = (320//8) * (320//8) + (320//16) * (320//16) + (320//32) * (320//32)
        
        outputs = np.random.randn(batch_size, total, channels).astype(np.float32)
        codeflash_output = demo_postprocess(outputs, img_size, p6=False); result = codeflash_output # 382μs -> 172μs (121% faster)

    def test_basic_with_p6_enabled(self):
        """Test basic functionality with p6=True configuration (8, 16, 32, 64 strides)."""
        img_size = (320, 320)
        batch_size = 1
        num_classes = 80
        channels = 85
        
        # Calculate expected total anchors with p6
        total = (320//8) * (320//8) + (320//16) * (320//16) + (320//32) * (320//32) + (320//64) * (320//64)
        
        outputs = np.random.randn(batch_size, total, channels).astype(np.float32)
        codeflash_output = demo_postprocess(outputs, img_size, p6=True); result = codeflash_output # 410μs -> 179μs (129% faster)

    def test_small_image_size(self):
        """Test with minimum viable image size of 64x64."""
        img_size = (64, 64)
        batch_size = 1
        channels = 85
        
        # Calculate anchors for 64x64 image
        total = (64//8) * (64//8) + (64//16) * (64//16) + (64//32) * (64//32)
        
        outputs = np.random.randn(batch_size, total, channels).astype(np.float32)
        codeflash_output = demo_postprocess(outputs, img_size, p6=False); result = codeflash_output # 206μs -> 58.3μs (255% faster)

    def test_output_values_are_transformed(self):
        """Test that output values are actually transformed (not returned unchanged)."""
        img_size = (64, 64)
        channels = 85
        total = (64//8) * (64//8) + (64//16) * (64//16) + (64//32) * (64//32)
        
        # Create outputs with specific values
        outputs = np.ones((1, total, channels), dtype=np.float32)
        original_outputs = outputs.copy()
        
        codeflash_output = demo_postprocess(outputs, img_size, p6=False); result = codeflash_output # 206μs -> 57.9μs (257% faster)

    def test_float32_dtype_preserved(self):
        """Test that float32 dtype is preserved in output."""
        img_size = (64, 64)
        channels = 85
        total = (64//8) * (64//8) + (64//16) * (64//16) + (64//32) * (64//32)
        
        outputs = np.random.randn(1, total, channels).astype(np.float32)
        codeflash_output = demo_postprocess(outputs, img_size, p6=False); result = codeflash_output # 207μs -> 57.8μs (260% faster)

    def test_float64_dtype_preserved(self):
        """Test that float64 dtype is preserved in output."""
        img_size = (64, 64)
        channels = 85
        total = (64//8) * (64//8) + (64//16) * (64//16) + (64//32) * (64//32)
        
        outputs = np.random.randn(1, total, channels).astype(np.float64)
        codeflash_output = demo_postprocess(outputs, img_size, p6=False); result = codeflash_output # 208μs -> 57.7μs (261% faster)

    def test_batch_size_greater_than_one(self):
        """Test processing multiple samples in batch."""
        img_size = (64, 64)
        batch_size = 4
        channels = 85
        total = (64//8) * (64//8) + (64//16) * (64//16) + (64//32) * (64//32)
        
        outputs = np.random.randn(batch_size, total, channels).astype(np.float32)
        codeflash_output = demo_postprocess(outputs, img_size, p6=False); result = codeflash_output # 228μs -> 84.9μs (169% faster)

    def test_width_height_channels_at_index_2_3(self):
        """Test that width and height transformations are applied at channels 2 and 3."""
        img_size = (64, 64)
        channels = 85
        total = (64//8) * (64//8) + (64//16) * (64//16) + (64//32) * (64//32)
        
        # Create outputs with specific values where we can track exp transformations
        outputs = np.zeros((1, total, channels), dtype=np.float32)
        outputs[:, :, 2] = 0.0  # width: exp(0) = 1
        outputs[:, :, 3] = 1.0  # height: exp(1) ≈ 2.718
        
        codeflash_output = demo_postprocess(outputs, img_size, p6=False); result = codeflash_output # 207μs -> 59.2μs (250% faster)

    def test_return_is_same_object(self):
        """Test that function returns the same array object (in-place modification)."""
        img_size = (64, 64)
        channels = 85
        total = (64//8) * (64//8) + (64//16) * (64//16) + (64//32) * (64//32)
        
        outputs = np.random.randn(1, total, channels).astype(np.float32)
        codeflash_output = demo_postprocess(outputs, img_size, p6=False); result = codeflash_output # 205μs -> 58.4μs (253% faster)

class TestDemoPostprocessEdgeCases:
    """Edge case test cases for unusual or extreme conditions."""
    
    def test_zero_values_in_outputs(self):
        """Test handling of zero values in the output tensor."""
        img_size = (64, 64)
        channels = 85
        total = (64//8) * (64//8) + (64//16) * (64//16) + (64//32) * (64//32)
        
        outputs = np.zeros((1, total, channels), dtype=np.float32)
        codeflash_output = demo_postprocess(outputs, img_size, p6=False); result = codeflash_output # 209μs -> 61.3μs (242% faster)

    def test_negative_values_in_outputs(self):
        """Test handling of negative values in the output tensor."""
        img_size = (64, 64)
        channels = 85
        total = (64//8) * (64//8) + (64//16) * (64//16) + (64//32) * (64//32)
        
        outputs = np.full((1, total, channels), -1.0, dtype=np.float32)
        codeflash_output = demo_postprocess(outputs, img_size, p6=False); result = codeflash_output # 204μs -> 59.2μs (246% faster)

    def test_large_positive_values_in_width_height(self):
        """Test handling of large positive values in width/height (channels 2-3)."""
        img_size = (64, 64)
        channels = 85
        total = (64//8) * (64//8) + (64//16) * (64//16) + (64//32) * (64//32)
        
        outputs = np.zeros((1, total, channels), dtype=np.float32)
        outputs[:, :, 2] = 10.0  # exp(10) is very large
        outputs[:, :, 3] = 10.0
        
        codeflash_output = demo_postprocess(outputs, img_size, p6=False); result = codeflash_output # 205μs -> 57.4μs (257% faster)

    def test_very_large_image_size(self):
        """Test with very large image size (1920x1920)."""
        img_size = (1920, 1920)
        channels = 85
        total = (1920//8) * (1920//8) + (1920//16) * (1920//16) + (1920//32) * (1920//32)
        
        outputs = np.random.randn(1, total, channels).astype(np.float32)
        codeflash_output = demo_postprocess(outputs, img_size, p6=False); result = codeflash_output # 10.2ms -> 7.98ms (27.3% faster)

    def test_non_square_image(self):
        """Test with non-square image dimensions."""
        img_size = (320, 640)  # width != height
        channels = 85
        total = (320//8) * (640//8) + (320//16) * (640//16) + (320//32) * (640//32)
        
        outputs = np.random.randn(1, total, channels).astype(np.float32)
        codeflash_output = demo_postprocess(outputs, img_size, p6=False); result = codeflash_output # 610μs -> 278μs (119% faster)

    def test_minimum_channels(self):
        """Test with minimum required channels (4 coordinates + 1 objectness)."""
        img_size = (64, 64)
        channels = 5  # minimum: x, y, w, h, objectness
        total = (64//8) * (64//8) + (64//16) * (64//16) + (64//32) * (64//32)
        
        outputs = np.random.randn(1, total, channels).astype(np.float32)
        codeflash_output = demo_postprocess(outputs, img_size, p6=False); result = codeflash_output # 224μs -> 73.9μs (204% faster)

    def test_many_classes(self):
        """Test with large number of classes (200 classes)."""
        img_size = (64, 64)
        channels = 5 + 200  # 4 coords + 1 objectness + 200 classes
        total = (64//8) * (64//8) + (64//16) * (64//16) + (64//32) * (64//32)
        
        outputs = np.random.randn(1, total, channels).astype(np.float32)
        codeflash_output = demo_postprocess(outputs, img_size, p6=False); result = codeflash_output # 213μs -> 60.5μs (253% faster)

    def test_nan_values_propagated(self):
        """Test that NaN values are propagated through the function."""
        img_size = (64, 64)
        channels = 85
        total = (64//8) * (64//8) + (64//16) * (64//16) + (64//32) * (64//32)
        
        outputs = np.ones((1, total, channels), dtype=np.float32)
        outputs[0, 0, 0] = np.nan
        
        codeflash_output = demo_postprocess(outputs, img_size, p6=False); result = codeflash_output # 208μs -> 57.7μs (261% faster)

    def test_inf_values_propagated(self):
        """Test that infinity values are propagated through the function."""
        img_size = (64, 64)
        channels = 85
        total = (64//8) * (64//8) + (64//16) * (64//16) + (64//32) * (64//32)
        
        outputs = np.ones((1, total, channels), dtype=np.float32)
        outputs[0, 0, 0] = np.inf
        
        codeflash_output = demo_postprocess(outputs, img_size, p6=False); result = codeflash_output # 205μs -> 58.4μs (252% faster)

    def test_p6_false_versus_true(self):
        """Test that p6=True creates different grid structures than p6=False."""
        img_size = (320, 320)
        channels = 85
        
        # For p6=False: strides are (8, 16, 32)
        total_no_p6 = (320//8) * (320//8) + (320//16) * (320//16) + (320//32) * (320//32)
        
        # For p6=True: strides are (8, 16, 32, 64)
        total_p6 = total_no_p6 + (320//64) * (320//64)
        
        outputs_no_p6 = np.random.randn(1, total_no_p6, channels).astype(np.float32)
        outputs_p6 = np.random.randn(1, total_p6, channels).astype(np.float32)
        
        codeflash_output = demo_postprocess(outputs_no_p6, img_size, p6=False); result_no_p6 = codeflash_output # 385μs -> 169μs (128% faster)
        codeflash_output = demo_postprocess(outputs_p6, img_size, p6=True); result_p6 = codeflash_output # 352μs -> 137μs (157% faster)

    def test_contiguous_and_non_contiguous_arrays(self):
        """Test that both contiguous and non-contiguous arrays are processed correctly."""
        img_size = (64, 64)
        channels = 85
        total = (64//8) * (64//8) + (64//16) * (64//16) + (64//32) * (64//32)
        
        # Create contiguous array
        outputs_contig = np.random.randn(1, total, channels).astype(np.float32)
        codeflash_output = demo_postprocess(outputs_contig, img_size, p6=False); result_contig = codeflash_output # 210μs -> 59.9μs (251% faster)
        
        # Create non-contiguous array by transposing and transposing back
        outputs_temp = np.random.randn(channels, total, 1).astype(np.float32)
        outputs_non_contig = np.transpose(outputs_temp, (2, 1, 0))
        codeflash_output = demo_postprocess(outputs_non_contig, img_size, p6=False); result_non_contig = codeflash_output # 171μs -> 39.9μs (329% faster)

class TestDemoPostprocessLargeScale:
    """Large scale test cases for performance and scalability assessment."""
    
    def test_large_batch_size(self):
        """Test with large batch size of 32 samples."""
        img_size = (320, 320)
        batch_size = 32
        channels = 85
        total = (320//8) * (320//8) + (320//16) * (320//16) + (320//32) * (320//32)
        
        outputs = np.random.randn(batch_size, total, channels).astype(np.float32)
        codeflash_output = demo_postprocess(outputs, img_size, p6=False); result = codeflash_output # 7.58ms -> 5.71ms (32.9% faster)

    def test_large_number_of_anchors(self):
        """Test with large number of anchor points (high resolution image 1280x1280)."""
        img_size = (1280, 1280)
        channels = 85
        total = (1280//8) * (1280//8) + (1280//16) * (1280//16) + (1280//32) * (1280//32)
        
        outputs = np.random.randn(1, total, channels).astype(np.float32)
        codeflash_output = demo_postprocess(outputs, img_size, p6=False); result = codeflash_output # 3.74ms -> 2.62ms (42.9% faster)

    def test_large_batch_and_image_size(self):
        """Test with both large batch size and large image size."""
        img_size = (960, 960)
        batch_size = 16
        channels = 85
        total = (960//8) * (960//8) + (960//16) * (960//16) + (960//32) * (960//32)
        
        outputs = np.random.randn(batch_size, total, channels).astype(np.float32)
        codeflash_output = demo_postprocess(outputs, img_size, p6=False); result = codeflash_output # 38.8ms -> 31.4ms (23.5% faster)

    def test_many_samples_sequential(self):
        """Test processing multiple images sequentially."""
        img_size = (320, 320)
        channels = 85
        total = (320//8) * (320//8) + (320//16) * (320//16) + (320//32) * (320//32)
        
        # Process 10 different outputs sequentially
        for _ in range(10):
            outputs = np.random.randn(1, total, channels).astype(np.float32)
            codeflash_output = demo_postprocess(outputs, img_size, p6=False); result = codeflash_output # 3.64ms -> 1.52ms (139% faster)

    def test_very_wide_aspect_ratio(self):
        """Test with extreme aspect ratio (1920x320)."""
        img_size = (1920, 320)
        channels = 85
        total = (1920//8) * (320//8) + (1920//16) * (320//16) + (1920//32) * (320//32)
        
        outputs = np.random.randn(1, total, channels).astype(np.float32)
        codeflash_output = demo_postprocess(outputs, img_size, p6=False); result = codeflash_output # 1.27ms -> 819μs (55.6% faster)

    def test_very_tall_aspect_ratio(self):
        """Test with extreme aspect ratio (320x1920)."""
        img_size = (320, 1920)
        channels = 85
        total = (320//8) * (1920//8) + (320//16) * (1920//16) + (320//32) * (1920//32)
        
        outputs = np.random.randn(1, total, channels).astype(np.float32)
        codeflash_output = demo_postprocess(outputs, img_size, p6=False); result = codeflash_output # 1.27ms -> 750μs (69.2% faster)

    def test_memory_efficiency_with_large_channels(self):
        """Test with very large channel dimension (1000 channels)."""
        img_size = (320, 320)
        channels = 1000
        total = (320//8) * (320//8) + (320//16) * (320//16) + (320//32) * (320//32)
        
        outputs = np.random.randn(1, total, channels).astype(np.float32)
        codeflash_output = demo_postprocess(outputs, img_size, p6=False); result = codeflash_output # 513μs -> 246μs (108% faster)

    def test_p6_with_large_image(self):
        """Test p6=True configuration with large image size."""
        img_size = (1024, 1024)
        channels = 85
        total = (1024//8) * (1024//8) + (1024//16) * (1024//16) + (1024//32) * (1024//32) + (1024//64) * (1024//64)
        
        outputs = np.random.randn(1, total, channels).astype(np.float32)
        codeflash_output = demo_postprocess(outputs, img_size, p6=True); result = codeflash_output # 2.66ms -> 1.47ms (81.4% faster)

    def test_stride_multiplication_correctness_large_scale(self):
        """Test that stride multiplication is correctly applied across all grid positions."""
        img_size = (256, 256)
        channels = 85
        total = (256//8) * (256//8) + (256//16) * (256//16) + (256//32) * (256//32)
        
        # Create output with known values
        outputs = np.zeros((1, total, channels), dtype=np.float32)
        outputs[:, :, 0] = 1.0  # x offset
        outputs[:, :, 1] = 1.0  # y offset
        
        codeflash_output = demo_postprocess(outputs, img_size, p6=False); result = codeflash_output # 324μs -> 131μs (147% faster)
        
        # Verify that different grid positions have different coordinates
        # due to grid addition and stride multiplication
        first_pos_x = result[0, 0, 0]
        last_pos_x = result[0, -1, 0]
# 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-demo_postprocess-mkovoubs and push.