⚡️ Speed up function compute_transpose_output_shape by 13%

[codeflash-ai]

📄 13% (0.13x) speedup for compute_transpose_output_shape in keras/src/ops/operation_utils.py

⏱️ Runtime : 288 microseconds → 256 microseconds (best of 250 runs)

📝 Explanation and details

The optimized code achieves a 12% speedup by eliminating an unnecessary list() conversion and leveraging tuple efficiency for indexing operations.

Key Optimizations:

1. Eliminated unnecessary list conversion: The original code converts input_shape to a list even when it's already a sequence. The optimized version converts directly to a tuple once, avoiding the intermediate list creation.

2. Leveraged tuple indexing efficiency: Tuples are more memory-efficient and faster for indexing operations than lists in Python. Since the function only needs to read values by index (not modify them), tuple is the optimal data structure.

3. Reduced memory allocations: By converting to tuple once and reusing it, the optimization reduces memory pressure and eliminates redundant conversions.

Performance Impact Analysis:

• Best gains on axes=None cases: Tests show 22-40% improvements when axes=None because tuple slicing (shape[::-1]) is significantly faster than list slicing
• Moderate gains on permutation cases: 5-13% improvements for custom axes due to faster tuple indexing
• Excellent scaling: Large shapes (1000 dimensions) see up to 171% improvement for reverse operations and 7-10% for permutations

Hot Path Impact:
Based on the function references, this optimization is particularly valuable because the function is called from:

• TensorFlow sparse tensor operations (tf.sparse.transpose)
• Keras tensor output shape computation in the ops layer

These are likely performance-critical paths where tensor shapes are computed frequently during model execution and compilation, making even small per-call improvements significant when aggregated across many operations.

The optimization maintains identical behavior and error handling while providing consistent performance gains across all test scenarios.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import pytest  # used for our unit tests
from keras.src.ops.operation_utils import compute_transpose_output_shape

# unit tests

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

def test_basic_reverse_axes_none():
    # Test with axes=None, should reverse the input shape
    codeflash_output = compute_transpose_output_shape((1, 2, 3), None) # 1.20μs -> 922ns (30.6% faster)

def test_basic_2d_swap():
    # Test with 2D shape and axes swapping
    codeflash_output = compute_transpose_output_shape((4, 5), [1, 0]) # 2.24μs -> 1.99μs (12.6% faster)

def test_basic_3d_permutation():
    # Test with 3D shape and a permutation of axes
    codeflash_output = compute_transpose_output_shape((2, 3, 4), [2, 0, 1]) # 2.01μs -> 1.90μs (5.62% faster)

def test_basic_identity_permutation():
    # Test with identity permutation (should return the same shape)
    codeflash_output = compute_transpose_output_shape((7, 8, 9), [0, 1, 2]) # 2.02μs -> 1.85μs (8.63% faster)

def test_basic_single_dim():
    # Test with a single dimension
    codeflash_output = compute_transpose_output_shape((42,), [0]) # 1.88μs -> 1.69μs (11.8% faster)

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

def test_edge_empty_shape():
    # Test with empty shape (no dimensions)
    codeflash_output = compute_transpose_output_shape((), None) # 1.19μs -> 937ns (26.9% faster)
    codeflash_output = compute_transpose_output_shape((), []) # 1.38μs -> 1.32μs (4.55% faster)

def test_edge_axes_length_mismatch_short():
    # Axes shorter than input_shape should raise ValueError
    with pytest.raises(ValueError):
        compute_transpose_output_shape((1, 2, 3), [0, 1]) # 1.76μs -> 1.76μs (0.284% faster)

def test_edge_axes_length_mismatch_long():
    # Axes longer than input_shape should raise ValueError
    with pytest.raises(ValueError):
        compute_transpose_output_shape((1, 2), [0, 1, 2]) # 1.75μs -> 1.72μs (1.86% faster)

def test_edge_axes_out_of_range():
    # Axes containing out-of-range indices should raise IndexError
    with pytest.raises(IndexError):
        compute_transpose_output_shape((1, 2, 3), [0, 1, 3]) # 2.55μs -> 2.50μs (1.84% faster)

def test_edge_axes_negative_index():
    # Axes containing negative indices should work (Python indexing)
    codeflash_output = compute_transpose_output_shape((10, 20, 30), [-1, 0, 1]) # 2.03μs -> 1.94μs (4.48% faster)

def test_edge_axes_repeated_index():
    # Axes containing repeated indices (not a valid permutation)
    # Should work but result in repeated dimensions
    codeflash_output = compute_transpose_output_shape((3, 4, 5), [1, 1, 2]) # 1.97μs -> 1.81μs (8.73% faster)

def test_edge_axes_unordered():
    # Axes in arbitrary order
    codeflash_output = compute_transpose_output_shape((1, 2, 3, 4), [3, 2, 1, 0]) # 2.03μs -> 1.80μs (12.8% faster)

def test_edge_axes_as_tuple():
    # Axes provided as tuple instead of list
    codeflash_output = compute_transpose_output_shape((1, 2, 3), (2, 0, 1)) # 1.95μs -> 1.75μs (11.3% faster)

def test_edge_axes_as_range():
    # Axes provided as a range object
    codeflash_output = compute_transpose_output_shape((5, 6, 7), range(3)) # 2.08μs -> 2.12μs (1.65% slower)

def test_edge_shape_with_zero():
    # Input shape contains zero
    codeflash_output = compute_transpose_output_shape((0, 4, 5), [2, 1, 0]) # 1.97μs -> 1.80μs (9.31% faster)

def test_edge_shape_with_negative():
    # Input shape contains negative numbers (unusual, but function should handle)
    codeflash_output = compute_transpose_output_shape((-1, 2, 3), [1, 2, 0]) # 1.90μs -> 1.76μs (7.89% faster)

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

def test_large_shape_reverse():
    # Large shape with axes=None (reverse)
    shape = tuple(range(1000))
    expected = tuple(reversed(shape))
    codeflash_output = compute_transpose_output_shape(shape, None) # 6.60μs -> 2.44μs (171% faster)

def test_large_shape_permutation():
    # Large shape with a shuffled permutation
    shape = tuple(range(1000))
    axes = list(range(999, -1, -1))  # reverse permutation
    expected = tuple(shape[ax] for ax in axes)
    codeflash_output = compute_transpose_output_shape(shape, axes) # 31.5μs -> 29.7μs (6.27% faster)

def test_large_shape_identity():
    # Large shape with identity permutation
    shape = tuple(range(1000))
    axes = list(range(1000))
    codeflash_output = compute_transpose_output_shape(shape, axes) # 32.1μs -> 29.8μs (7.71% faster)

def test_large_shape_random_permutation():
    # Large shape with a random permutation
    import random
    shape = tuple(range(1000))
    axes = list(range(1000))
    random.seed(42)
    random.shuffle(axes)
    expected = tuple(shape[ax] for ax in axes)
    codeflash_output = compute_transpose_output_shape(shape, axes) # 32.2μs -> 30.1μs (7.08% faster)

def test_large_shape_axes_length_mismatch():
    # Large shape with axes length mismatch
    shape = tuple(range(1000))
    axes = list(range(999))  # one less than shape
    with pytest.raises(ValueError):
        compute_transpose_output_shape(shape, axes) # 3.56μs -> 2.00μs (77.5% faster)

def test_large_shape_axes_out_of_range():
    # Large shape with axes containing out-of-range index
    shape = tuple(range(1000))
    axes = list(range(1000))
    axes[0] = 1000  # out of range
    with pytest.raises(IndexError):
        compute_transpose_output_shape(shape, axes) # 4.11μs -> 2.56μs (60.4% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import pytest  # used for our unit tests
from keras.src.ops.operation_utils import compute_transpose_output_shape

# unit tests

# --- Basic Test Cases ---

def test_basic_2d_reverse_axes():
    # 2D shape, axes=None should reverse dimensions
    codeflash_output = compute_transpose_output_shape((2, 3), None) # 1.45μs -> 1.03μs (40.1% faster)

def test_basic_3d_reverse_axes():
    # 3D shape, axes=None should reverse dimensions
    codeflash_output = compute_transpose_output_shape((2, 3, 4), None) # 1.24μs -> 910ns (36.2% faster)

def test_basic_2d_custom_axes():
    # 2D shape, axes=[1,0] should swap dimensions
    codeflash_output = compute_transpose_output_shape((2, 3), [1, 0]) # 2.17μs -> 1.97μs (10.2% faster)

def test_basic_3d_custom_axes():
    # 3D shape, axes=[2,0,1] should permute dimensions
    codeflash_output = compute_transpose_output_shape((2, 3, 4), [2, 0, 1]) # 2.04μs -> 1.88μs (8.52% faster)

def test_basic_4d_identity_axes():
    # 4D shape, axes=[0,1,2,3] (identity) should return input shape
    codeflash_output = compute_transpose_output_shape((2, 3, 4, 5), [0, 1, 2, 3]) # 2.02μs -> 1.91μs (5.55% faster)

def test_basic_4d_reverse_axes():
    # 4D shape, axes=None should reverse dimensions
    codeflash_output = compute_transpose_output_shape((2, 3, 4, 5), None) # 1.27μs -> 917ns (38.9% faster)

# --- Edge Test Cases ---

def test_edge_empty_shape():
    # Empty input shape, should return empty tuple
    codeflash_output = compute_transpose_output_shape((), None) # 1.10μs -> 895ns (22.9% faster)
    codeflash_output = compute_transpose_output_shape((), []) # 1.49μs -> 1.45μs (3.18% faster)

def test_edge_singleton_shape():
    # Singleton shape, axes=None should return same shape
    codeflash_output = compute_transpose_output_shape((7,), None) # 1.13μs -> 808ns (39.5% faster)
    # Singleton shape, axes=[0] should return same shape
    codeflash_output = compute_transpose_output_shape((7,), [0]) # 1.65μs -> 1.49μs (10.9% faster)

def test_edge_axes_length_mismatch_shorter():
    # Axes shorter than input_shape should raise ValueError
    with pytest.raises(ValueError):
        compute_transpose_output_shape((2, 3), [0]) # 1.84μs -> 1.72μs (7.28% faster)

def test_edge_axes_length_mismatch_longer():
    # Axes longer than input_shape should raise ValueError
    with pytest.raises(ValueError):
        compute_transpose_output_shape((2, 3), [0, 1, 2]) # 1.70μs -> 1.73μs (1.28% slower)

def test_edge_axes_out_of_bounds():
    # Axes with out-of-bounds index should raise IndexError
    with pytest.raises(IndexError):
        compute_transpose_output_shape((2, 3), [0, 2]) # 2.45μs -> 2.53μs (3.08% slower)

def test_edge_axes_negative_index():
    # Axes with negative index should work like Python indexing
    codeflash_output = compute_transpose_output_shape((2, 3, 4), [0, -1, 1]) # 2.05μs -> 1.94μs (5.50% faster)

def test_edge_axes_with_duplicates():
    # Axes with duplicate indices should work, but may not be meaningful
    codeflash_output = compute_transpose_output_shape((2, 3), [0, 0]) # 1.96μs -> 1.77μs (10.6% faster)

def test_edge_axes_with_repeated_negative_index():
    # Negative indices repeated
    codeflash_output = compute_transpose_output_shape((2, 3, 4), [-1, -1, -1]) # 1.99μs -> 1.72μs (15.7% faster)

def test_edge_shape_with_zeros():
    # Input shape with zeros, should preserve zeros
    codeflash_output = compute_transpose_output_shape((0, 2, 3), [2, 1, 0]) # 2.02μs -> 1.78μs (13.7% faster)

def test_edge_shape_with_negative_dim():
    # Input shape with negative dimension, should preserve negative
    codeflash_output = compute_transpose_output_shape((-2, 3, 4), [2, 0, 1]) # 1.97μs -> 1.76μs (11.8% faster)

# --- Large Scale Test Cases ---

def test_large_scale_shape_and_axes():
    # Large shape and axes, 1000 elements, axes=None (reverse)
    shape = tuple(range(1000))
    expected = tuple(reversed(shape))
    codeflash_output = compute_transpose_output_shape(shape, None) # 6.59μs -> 2.50μs (163% faster)

def test_large_scale_custom_axes():
    # Large shape and axes, 1000 elements, custom permutation
    shape = tuple(range(1000))
    axes = list(reversed(range(1000)))
    expected = tuple(shape[ax] for ax in axes)
    codeflash_output = compute_transpose_output_shape(shape, axes) # 32.2μs -> 29.4μs (9.64% faster)

def test_large_scale_identity_axes():
    # Large shape and axes, 1000 elements, identity permutation
    shape = tuple(range(1000))
    axes = list(range(1000))
    codeflash_output = compute_transpose_output_shape(shape, axes) # 32.4μs -> 30.2μs (7.03% faster)

def test_large_scale_axes_length_mismatch():
    # Large shape, axes length mismatch
    shape = tuple(range(1000))
    axes = list(range(999))
    with pytest.raises(ValueError):
        compute_transpose_output_shape(shape, axes) # 3.53μs -> 2.00μs (76.7% faster)

def test_large_scale_axes_out_of_bounds():
    # Large shape, axes out of bounds
    shape = tuple(range(1000))
    axes = list(range(999)) + [1001]
    with pytest.raises(IndexError):
        compute_transpose_output_shape(shape, axes) # 33.8μs -> 32.1μs (5.52% faster)

# --- Additional Edge Cases ---

def test_axes_is_empty_list_with_nonempty_shape():
    # Axes is empty list but shape is not, should raise ValueError
    with pytest.raises(ValueError):
        compute_transpose_output_shape((1,), []) # 1.81μs -> 1.77μs (1.92% faster)

def test_axes_is_none_with_non_tuple_shape():
    # Input shape is list, axes=None, should reverse
    codeflash_output = compute_transpose_output_shape([1, 2, 3], None) # 1.40μs -> 1.11μs (25.3% faster)

def test_axes_is_none_with_string_shape():
    # Input shape is string, which is iterable, should reverse
    codeflash_output = compute_transpose_output_shape("abc", None) # 1.39μs -> 1.23μs (13.4% faster)

def test_axes_is_none_with_range_shape():
    # Input shape is range, should reverse
    codeflash_output = compute_transpose_output_shape(range(3), None) # 1.50μs -> 1.31μs (14.0% 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-compute_transpose_output_shape-mjagndd1 and push.