⚡️ Speed up function compute_pooling_output_shape by 10%

[codeflash-ai]

📄 10% (0.10x) speedup for compute_pooling_output_shape in keras/src/ops/operation_utils.py

⏱️ Runtime : 666 microseconds → 604 microseconds (best of 23 runs)

📝 Explanation and details

The optimization eliminates unnecessary NumPy array creation by removing the conversion of the entire input_shape to a NumPy array early in the function, and instead converts only the required variables (spatial_shape, pool_size, and strides) to NumPy arrays when needed for computation.

Key changes:

• Removed input_shape = np.array(input_shape) which was creating an unnecessary array copy of the entire input shape
• Added spatial_shape = np.array(spatial_shape) and strides = np.array(strides) right before the mathematical operations that require NumPy arrays
• This reduces memory allocation overhead and avoids converting dimensions that aren't used in calculations

Why this is faster:
The original code converted the entire input shape to a NumPy array (15.8% of total time) but only needed the spatial dimensions as arrays for the pooling calculations. The optimization defers array creation until actually needed and only converts the minimal required data structures. This reduces both memory allocation overhead and CPU cycles spent on unnecessary data type conversions.

Impact on workloads:
Based on the function references, this function is called during layer initialization and tensor shape computation in pooling layers, which happens during model construction and inference. The 10% speedup will benefit:

• Model compilation time when pooling layers compute output shapes
• Dynamic shape inference during training/inference
• Particularly beneficial for models with many pooling layers or frequent shape recomputation

Test case performance:
The optimization shows consistent 7-17% improvements across various test cases, with larger gains on edge cases involving stride calculations and different data formats, indicating the optimization is broadly effective across different pooling configurations.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import numpy as np
# imports
import pytest
from keras.src.ops.operation_utils import compute_pooling_output_shape

# unit tests

# --- Basic Test Cases ---

def test_basic_strided_pooling():
    # Pooling with strides different from pool_size
    codeflash_output = compute_pooling_output_shape((1, 4, 4, 1), (2, 2), strides=(1, 1)) # 35.7μs -> 33.0μs (8.20% faster)

def test_strides_none_defaults_to_pool_size():
    # If strides is None, should default to pool_size
    codeflash_output = compute_pooling_output_shape((1, 8, 8, 1), (2, 2), strides=None) # 36.3μs -> 33.8μs (7.43% faster)

def test_stride_larger_than_input():
    # Stride larger than input, should result in 1 output spatial location
    codeflash_output = compute_pooling_output_shape((1, 2, 2, 1), (1, 1), strides=(3, 3)) # 41.1μs -> 35.1μs (17.2% faster)

def test_non_integer_output_shape():
    # Output shape should be floored
    codeflash_output = compute_pooling_output_shape((1, 7, 7, 1), (3, 3), strides=(2, 2)) # 23.3μs -> 21.6μs (7.87% faster)

def test_large_stride():
    # Large stride, output spatial dims should be small
    codeflash_output = compute_pooling_output_shape((1, 100, 100, 1), (3, 3), strides=(50, 50)) # 39.1μs -> 34.9μs (11.9% faster)

def test_strides_not_tuple():
    # Strides as list instead of tuple should still work
    codeflash_output = compute_pooling_output_shape((1, 8, 8, 1), (2, 2), strides=[2, 2]) # 38.0μs -> 33.7μs (12.9% faster)

import numpy as np
# imports
import pytest
from keras.src.ops.operation_utils import compute_pooling_output_shape

# unit tests

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

def test_basic_strided_pooling():
    # 4x4 input, 2x2 pooling, stride 1, batch=1, channels=1
    codeflash_output = compute_pooling_output_shape((1, 4, 4, 1), (2, 2), strides=(1, 1)) # 39.4μs -> 34.3μs (14.9% faster)

def test_basic_1d_pooling():
    # 1D pooling: (batch, steps, channels)
    codeflash_output = compute_pooling_output_shape((8, 10, 3), (2,), strides=(2,)) # 37.2μs -> 33.4μs (11.3% faster)

def test_basic_same_padding_stride1():
    # 4x4 input, 2x2 pooling, stride 1, 'same' padding
    codeflash_output = compute_pooling_output_shape((1, 4, 4, 1), (2, 2), strides=(1, 1), padding='same') # 37.1μs -> 31.6μs (17.5% faster)

def test_basic_different_pool_and_stride():
    # 5x5 input, 2x2 pooling, 3x3 strides
    codeflash_output = compute_pooling_output_shape((1, 5, 5, 1), (2, 2), strides=(3, 3)) # 23.6μs -> 22.1μs (6.77% faster)

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

def test_edge_stride_larger_than_input():
    # Stride > input size, should return 1x1
    codeflash_output = compute_pooling_output_shape((1, 2, 2, 1), (1, 1), strides=(3, 3)) # 38.0μs -> 33.7μs (12.6% faster)

def test_edge_invalid_data_format():
    # Invalid data format should raise IndexError or ValueError, depending on implementation
    with pytest.raises(Exception):
        compute_pooling_output_shape((1, 4, 4, 1), (2, 2), data_format='NCHW') # 3.73μs -> 3.74μs (0.160% slower)

def test_edge_stride_none():
    # Stride=None should default to pool_size
    codeflash_output = compute_pooling_output_shape((1, 6, 6, 1), (3, 3), strides=None) # 36.8μs -> 32.1μs (14.4% faster)

def test_edge_non_integer_input_shape():
    # Non-integer input shape should raise or fail
    with pytest.raises(Exception):
        compute_pooling_output_shape((1, 4.5, 4, 1), (2, 2)) # 2.47μs -> 2.57μs (4.16% slower)

def test_edge_negative_input_shape():
    # Negative input shape should raise or produce negative output
    with pytest.raises(Exception):
        compute_pooling_output_shape((1, -4, 4, 1), (2, 2)) # 3.55μs -> 3.38μs (4.82% faster)

def test_edge_pool_size_zero():
    # Pool size zero should raise or produce error
    with pytest.raises(Exception):
        compute_pooling_output_shape((1, 4, 4, 1), (0, 2)) # 2.69μs -> 2.51μs (7.12% faster)

def test_edge_stride_zero():
    # Stride zero should raise or produce error (division by zero)
    with pytest.raises(Exception):
        compute_pooling_output_shape((1, 4, 4, 1), (2, 2), strides=(0, 1)) # 43.5μs -> 40.1μs (8.58% faster)

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

def test_large_stride():
    # Large stride, should reduce output size drastically
    codeflash_output = compute_pooling_output_shape((1, 100, 100, 1), (2, 2), strides=(50, 50)) # 37.8μs -> 34.4μs (9.84% faster)

def test_large_dim_channels_first():
    # Large dims with channels_first
    codeflash_output = compute_pooling_output_shape((10, 3, 128, 128), (8, 8), strides=(8, 8), data_format="channels_first") # 36.0μs -> 33.2μs (8.70% faster)

# -------------------- Parameterized Tests for Scalability --------------------

@pytest.mark.parametrize("batch, h, w, pool, stride, expected", [
    (1, 100, 100, (10, 10), (10, 10), (1, 10, 10, 1)),
    (8, 128, 64, (16, 8), (16, 8), (8, 8, 8, 1)),
    (32, 256, 256, (32, 32), (32, 32), (32, 8, 8, 1)),
])
def test_param_large_cases(batch, h, w, pool, stride, expected):
    # Parameterized large cases for coverage
    codeflash_output = compute_pooling_output_shape((batch, h, w, 1), pool, strides=stride) # 81.1μs -> 75.1μs (7.94% faster)

# -------------------- Mutation Testing Guard --------------------

def test_mutation_guard():
    # Changing any core arithmetic should fail this test
    # 7x7 input, 3x3 pooling, stride 2, 'valid' padding
    codeflash_output = compute_pooling_output_shape((1, 7, 7, 1), (3, 3), strides=(2, 2)) # 21.6μs -> 21.6μs (0.111% slower)
    # 7x7 input, 3x3 pooling, stride 2, 'same' padding
    codeflash_output = compute_pooling_output_shape((1, 7, 7, 1), (3, 3), strides=(2, 2), padding='same') # 8.30μs -> 8.31μs (0.132% slower)
# 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-compute_pooling_output_shape-mjafzodd and push.