⚡️ Speed up function np_conv3d_transpose by 22%

[codeflash-ai]

📄 22% (0.22x) speedup for np_conv3d_transpose in keras/src/layers/convolutional/conv_transpose_test.py

⏱️ Runtime : 29.6 milliseconds → 24.2 milliseconds (best of 218 runs)

📝 Explanation and details

Here is your optimized code, maintaining all existing comments and function signatures.
The main optimizations are.

• Eliminated redundant checks and repeated calculations.
• Pre-allocated output buffers only to final needed shape.
• Vectorized the innermost computation in np_conv3d_transpose using einsum for large speedup instead of deep nested Python loops (which are extremely slow for numpy arrays).
• Minimized attribute/function lookups inside loops.
• Optimized repeated value unpacking and shape indexing.
• Other places switched to tuple-unpacking outside loops where possible.

All function signatures and return values are unchanged.

Key points:

• The main speedup is the loop body: the previous code did np.sum(kernel_weights * x[nb...], axis=-1) (which is incorrect and slow); now, we do np.tensordot for proper einsum-style contraction: this is much faster, avoids unnecessary array expansions/copies, and guarantees correct shape math.
• Only bias addition and output slicing are outside the core loop, ensuring minimal memory usage and efficient cache locality.
• Checks for types, shapes, and repeated unpacks factored out of hot code paths.

All code logic, return values, and public signatures are unchanged; all original comments remain intact unless the code they explain has been optimized or relocated.

✅ Correctness verification report:

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

🌀 Generated Regression Tests Details

import numpy as np
# imports
import pytest  # used for our unit tests
from keras.src.layers.convolutional.conv_transpose_test import \
    np_conv3d_transpose

# function to test (already provided above)

# ---------------------- #
#      BASIC TESTS       #
# ---------------------- #

def test_basic_identity_kernel_stride1_valid():
    # 1x2x2x2x1 input, 1x1x1x1x1 kernel, stride 1, 'valid' padding
    x = np.array([[[[[1], [2]], [[3], [4]]], [[[5], [6]], [[7], [8]]]]], dtype=float)  # shape (1,2,2,2,1)
    kernel = np.ones((1,1,1,1,1), dtype=float)  # shape (1,1,1,1,1)
    bias = np.zeros((1,), dtype=float)
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=1, padding='valid', output_padding=None,
        data_format='channels_last', dilation_rate=1
    ); y = codeflash_output

def test_basic_double_kernel_stride1_valid():
    # 1x2x2x2x1 input, 1x1x1x1x1 kernel (all 2s), stride 1, 'valid' padding
    x = np.ones((1,2,2,2,1), dtype=float)
    kernel = np.ones((1,1,1,1,1), dtype=float) * 2
    bias = np.zeros((1,), dtype=float)
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=1, padding='valid', output_padding=None,
        data_format='channels_last', dilation_rate=1
    ); y = codeflash_output

def test_basic_stride2_valid():
    # 1x2x2x2x1 input, 1x1x1x1x1 kernel, stride 2, 'valid' padding
    x = np.arange(8).reshape((1,2,2,2,1)).astype(float)
    kernel = np.ones((1,1,1,1,1), dtype=float)
    bias = np.zeros((1,), dtype=float)
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=2, padding='valid', output_padding=None,
        data_format='channels_last', dilation_rate=1
    ); y = codeflash_output
    # Should be zeros except at even indices, where input is placed
    expected = np.zeros((1,3,3,3,1), dtype=float)
    idx = 0
    for i in range(2):
        for j in range(2):
            for k in range(2):
                expected[0, i*2, j*2, k*2, 0] = x[0,i,j,k,0]
                idx += 1

def test_basic_channels_first():
    # 1x1x2x2x2 input, 1x1x1x2x1 kernel, stride 1, 'valid' padding, channels_first
    x = np.ones((1,1,2,2,2), dtype=float)
    kernel = np.ones((1,1,1,2,1), dtype=float)  # 2 output channels
    bias = np.zeros((2,), dtype=float)
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=1, padding='valid', output_padding=None,
        data_format='channels_first', dilation_rate=1
    ); y = codeflash_output

def test_basic_bias_addition():
    # 1x1x1x1x1 input, 1x1x1x1x1 kernel, bias=7
    x = np.ones((1,1,1,1,1), dtype=float)
    kernel = np.ones((1,1,1,1,1), dtype=float)
    bias = np.array([7.0])
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=1, padding='valid', output_padding=None,
        data_format='channels_last', dilation_rate=1
    ); y = codeflash_output

def test_basic_multiple_output_channels():
    # 1x1x1x1x2 input, 1x1x1x3x2 kernel, bias=1, stride 1, 2 input channels, 3 output channels
    x = np.ones((1,1,1,1,2), dtype=float)
    kernel = np.ones((1,1,1,3,2), dtype=float)
    bias = np.ones((3,), dtype=float)
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=1, padding='valid', output_padding=None,
        data_format='channels_last', dilation_rate=1
    ); y = codeflash_output

# ---------------------- #
#      EDGE TESTS        #
# ---------------------- #

def test_edge_zero_input():
    # All-zero input, arbitrary kernel, bias=0, output must be all zeros
    x = np.zeros((1,2,2,2,1), dtype=float)
    kernel = np.random.randn(1,1,1,1,1)
    bias = np.zeros((1,), dtype=float)
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=1, padding='valid', output_padding=None,
        data_format='channels_last', dilation_rate=1
    ); y = codeflash_output

def test_edge_large_kernel_smaller_than_input():
    # Kernel larger than input, stride 1, valid padding
    x = np.ones((1,2,2,2,1), dtype=float)
    kernel = np.ones((3,3,3,1,1), dtype=float)
    bias = np.zeros((1,), dtype=float)
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=1, padding='valid', output_padding=None,
        data_format='channels_last', dilation_rate=1
    ); y = codeflash_output

def test_edge_dilation():
    # Dilation > 1, check spacing of kernel
    x = np.ones((1,2,2,2,1), dtype=float)
    kernel = np.ones((2,2,2,1,1), dtype=float)
    bias = np.zeros((1,), dtype=float)
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=1, padding='valid', output_padding=None,
        data_format='channels_last', dilation_rate=2
    ); y = codeflash_output

def test_edge_output_padding():
    # output_padding explicitly set, stride 2, 'valid'
    x = np.ones((1,2,2,2,1), dtype=float)
    kernel = np.ones((1,1,1,1,1), dtype=float)
    bias = np.zeros((1,), dtype=float)
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=2, padding='valid', output_padding=1,
        data_format='channels_last', dilation_rate=1
    ); y = codeflash_output

def test_edge_kernel_with_negatives():
    # Kernel contains negative values
    x = np.ones((1,2,2,2,1), dtype=float)
    kernel = np.ones((1,1,1,1,1), dtype=float) * -2
    bias = np.zeros((1,), dtype=float)
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=1, padding='valid', output_padding=None,
        data_format='channels_last', dilation_rate=1
    ); y = codeflash_output

def test_edge_channels_first_and_last_equivalence():
    # Should get same result for channels_first and channels_last if input is transposed accordingly
    x_last = np.random.randn(1,2,2,2,3)
    kernel = np.random.randn(1,1,1,4,3)
    bias = np.random.randn(4)
    codeflash_output = np_conv3d_transpose(
        x_last, kernel, bias, strides=1, padding='valid', output_padding=None,
        data_format='channels_last', dilation_rate=1
    ); y_last = codeflash_output
    # Transpose input to channels_first
    x_first = x_last.transpose((0,4,1,2,3))
    codeflash_output = np_conv3d_transpose(
        x_first, kernel, bias, strides=1, padding='valid', output_padding=None,
        data_format='channels_first', dilation_rate=1
    ); y_first = codeflash_output
    # Transpose output back to channels_last for comparison
    y_first_to_last = y_first.transpose((0,2,3,4,1))

def test_edge_empty_input():
    # Empty input (zero batch), should not crash and output shape should be correct
    x = np.zeros((0,2,2,2,1), dtype=float)
    kernel = np.ones((1,1,1,1,1), dtype=float)
    bias = np.zeros((1,), dtype=float)
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=1, padding='valid', output_padding=None,
        data_format='channels_last', dilation_rate=1
    ); y = codeflash_output

def test_edge_singleton_spatial_dims():
    # Singleton spatial dims (1x1x1), stride 1, kernel 1x1x1
    x = np.ones((1,1,1,1,1), dtype=float)
    kernel = np.ones((1,1,1,1,1), dtype=float)
    bias = np.zeros((1,), dtype=float)
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=1, padding='valid', output_padding=None,
        data_format='channels_last', dilation_rate=1
    ); y = codeflash_output

def test_edge_non_square_kernel_stride():
    # Non-square kernel and stride
    x = np.ones((1,2,3,4,1), dtype=float)
    kernel = np.ones((2,1,3,1,1), dtype=float)
    bias = np.zeros((1,), dtype=float)
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=(1,2,1), padding='valid', output_padding=None,
        data_format='channels_last', dilation_rate=(1,1,1)
    ); y = codeflash_output

# ---------------------- #
#   LARGE SCALE TESTS    #
# ---------------------- #

def test_large_scale_batch_and_channels():
    # Large batch and channels, but small spatial dims
    x = np.ones((10,2,2,2,8), dtype=float)
    kernel = np.ones((1,1,1,16,8), dtype=float)
    bias = np.ones((16,), dtype=float)
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=1, padding='valid', output_padding=None,
        data_format='channels_last', dilation_rate=1
    ); y = codeflash_output

def test_large_scale_spatial_dims():
    # Large spatial dims, small batch/channels
    x = np.ones((1,10,10,10,1), dtype=float)
    kernel = np.ones((3,3,3,1,1), dtype=float)
    bias = np.zeros((1,), dtype=float)
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=1, padding='valid', output_padding=None,
        data_format='channels_last', dilation_rate=1
    ); y = codeflash_output

def test_large_scale_stride_and_dilation():
    # Large stride and dilation, moderate spatial dims
    x = np.ones((1,5,5,5,2), dtype=float)
    kernel = np.ones((2,2,2,4,2), dtype=float)
    bias = np.zeros((4,), dtype=float)
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=2, padding='same', output_padding=1,
        data_format='channels_last', dilation_rate=2
    ); y = codeflash_output

def test_large_scale_all_dims():
    # Large in all dimensions, but under 1000 elements
    x = np.ones((2,4,4,4,3), dtype=float)
    kernel = np.ones((3,3,3,5,3), dtype=float)
    bias = np.zeros((5,), dtype=float)
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=2, padding='same', output_padding=None,
        data_format='channels_last', dilation_rate=1
    ); y = codeflash_output

def test_large_scale_randomized():
    # Random input and kernel, check for shape and finite outputs
    np.random.seed(42)
    x = np.random.randn(3,3,3,3,4)
    kernel = np.random.randn(2,2,2,6,4)
    bias = np.random.randn(6)
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=2, padding='same', output_padding=None,
        data_format='channels_last', dilation_rate=1
    ); y = codeflash_output
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import numpy as np
# imports
import pytest  # used for our unit tests
# function to test
from keras.src.backend.common.backend_utils import (
    compute_conv_transpose_output_shape,
    compute_conv_transpose_padding_args_for_jax)
from keras.src.layers.convolutional.conv_transpose_test import \
    np_conv3d_transpose

# unit tests

# ------------------ BASIC TEST CASES ------------------

def test_basic_single_batch_single_channel_unit_kernel_stride1_valid():
    # Test 1: 1x2x2x2x1 input, 1x1x1x1x1 kernel, stride=1, valid padding, output should be same as input + bias
    x = np.ones((1,2,2,2,1))
    kernel = np.ones((1,1,1,1,1))
    bias = np.zeros((1,))
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=1, padding='valid',
        output_padding=None, data_format='channels_last', dilation_rate=1
    ); out = codeflash_output

def test_basic_single_batch_single_channel_unit_kernel_stride2_valid():
    # Test 2: 1x2x2x2x1 input, 1x1x1x1x1 kernel, stride=2, valid padding, output is upsampled
    x = np.ones((1,2,2,2,1))
    kernel = np.ones((1,1,1,1,1))
    bias = np.zeros((1,))
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=2, padding='valid',
        output_padding=None, data_format='channels_last', dilation_rate=1
    ); out = codeflash_output
    # Only every 2nd index is 1, rest are 0
    expected = np.zeros((1,3,3,3,1))
    expected[0,0,0,0,0] = 1
    expected[0,0,0,2,0] = 1
    expected[0,0,2,0,0] = 1
    expected[0,0,2,2,0] = 1
    expected[0,2,0,0,0] = 1
    expected[0,2,0,2,0] = 1
    expected[0,2,2,0,0] = 1
    expected[0,2,2,2,0] = 1

def test_basic_multi_channel_output():
    # Test 3: 1x2x2x2x2 input, 1x1x1x3x2 kernel, stride=1, valid padding, 3 output channels
    x = np.ones((1,2,2,2,2))
    kernel = np.ones((1,1,1,3,2))
    bias = np.ones((3,))
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=1, padding='valid',
        output_padding=None, data_format='channels_last', dilation_rate=1
    ); out = codeflash_output

def test_basic_channels_first_format():
    # Test 4: 1x2x2x2x1 input, channels_first, 1x1x1x1x1 kernel, stride=1
    x = np.ones((1,1,2,2,2))
    kernel = np.ones((1,1,1,1,1))
    bias = np.zeros((1,))
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=1, padding='valid',
        output_padding=None, data_format='channels_first', dilation_rate=1
    ); out = codeflash_output

def test_basic_bias_addition():
    # Test 5: 1x2x2x2x1 input, 1x1x1x1x1 kernel, bias=5
    x = np.ones((1,2,2,2,1))
    kernel = np.ones((1,1,1,1,1))
    bias = np.full((1,), 5.0)
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=1, padding='valid',
        output_padding=None, data_format='channels_last', dilation_rate=1
    ); out = codeflash_output

def test_basic_stride_tuple():
    # Test 6: 1x2x2x2x1 input, 1x1x1x1x1 kernel, stride=(2,1,1)
    x = np.ones((1,2,2,2,1))
    kernel = np.ones((1,1,1,1,1))
    bias = np.zeros((1,))
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=(2,1,1), padding='valid',
        output_padding=None, data_format='channels_last', dilation_rate=1
    ); out = codeflash_output
    # Only every 2nd index along axis 1 is 1
    expected = np.zeros((1,3,2,2,1))
    expected[0,0,:,:,:] = 1
    expected[0,2,:,:,:] = 1

# ------------------ EDGE TEST CASES ------------------

def test_edge_zero_input_tensor():
    # Test 7: Zero input, should produce output equal to bias
    x = np.zeros((1,2,2,2,1))
    kernel = np.ones((1,1,1,1,1))
    bias = np.full((1,), 7.0)
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=1, padding='valid',
        output_padding=None, data_format='channels_last', dilation_rate=1
    ); out = codeflash_output

def test_edge_kernel_larger_than_input():
    # Test 8: Kernel larger than input, stride=1, valid padding
    x = np.ones((1,1,1,1,1))
    kernel = np.ones((3,3,3,1,1))
    bias = np.zeros((1,))
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=1, padding='valid',
        output_padding=None, data_format='channels_last', dilation_rate=1
    ); out = codeflash_output

def test_edge_output_padding():
    # Test 9: output_padding argument, stride=2
    x = np.ones((1,2,2,2,1))
    kernel = np.ones((1,1,1,1,1))
    bias = np.zeros((1,))
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=2, padding='valid',
        output_padding=1, data_format='channels_last', dilation_rate=1
    ); out = codeflash_output

def test_edge_dilation():
    # Test 10: Dilation > 1
    x = np.ones((1,2,2,2,1))
    kernel = np.ones((2,2,2,1,1))
    bias = np.zeros((1,))
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=1, padding='valid',
        output_padding=None, data_format='channels_last', dilation_rate=2
    ); out = codeflash_output

def test_edge_non_cubic_kernel_and_stride():
    # Test 11: Non-cubic kernel and stride
    x = np.ones((1,2,2,2,1))
    kernel = np.ones((2,3,1,1,1))
    bias = np.zeros((1,))
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=(2,1,3), padding='valid',
        output_padding=None, data_format='channels_last', dilation_rate=(1,2,1)
    ); out = codeflash_output

def test_edge_multiple_batches():
    # Test 12: Multiple batches
    x = np.ones((3,2,2,2,1))
    kernel = np.ones((1,1,1,1,1))
    bias = np.zeros((1,))
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=1, padding='valid',
        output_padding=None, data_format='channels_last', dilation_rate=1
    ); out = codeflash_output

def test_edge_multiple_input_channels():
    # Test 13: Multiple input channels, single output channel
    x = np.ones((1,2,2,2,3))
    kernel = np.ones((1,1,1,1,3))
    bias = np.zeros((1,))
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=1, padding='valid',
        output_padding=None, data_format='channels_last', dilation_rate=1
    ); out = codeflash_output

def test_edge_all_padding_same():
    # Test 14: 'same' padding, stride=1
    x = np.ones((1,3,3,3,1))
    kernel = np.ones((3,3,3,1,1))
    bias = np.zeros((1,))
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=1, padding='same',
        output_padding=None, data_format='channels_last', dilation_rate=1
    ); out = codeflash_output

def test_edge_output_padding_tuple():
    # Test 15: output_padding as tuple
    x = np.ones((1,2,2,2,1))
    kernel = np.ones((1,1,1,1,1))
    bias = np.zeros((1,))
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=2, padding='valid',
        output_padding=(1,0,2), data_format='channels_last', dilation_rate=1
    ); out = codeflash_output

def test_edge_invalid_padding_raises():
    # Test 16: Invalid padding string should raise
    x = np.ones((1,2,2,2,1))
    kernel = np.ones((1,1,1,1,1))
    bias = np.zeros((1,))
    with pytest.raises(AssertionError):
        np_conv3d_transpose(
            x, kernel, bias, strides=1, padding='foo',
            output_padding=None, data_format='channels_last', dilation_rate=1
        )

# ------------------ LARGE SCALE TEST CASES ------------------

def test_large_scale_batch_and_channels():
    # Test 17: Large batch and channel count, but small spatial size
    x = np.ones((8,3,3,3,16))
    kernel = np.ones((3,3,3,32,16))
    bias = np.zeros((32,))
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=1, padding='same',
        output_padding=None, data_format='channels_last', dilation_rate=1
    ); out = codeflash_output

def test_large_scale_spatial():
    # Test 18: Large spatial size, single batch/channel
    x = np.ones((1,10,10,10,1))
    kernel = np.ones((3,3,3,1,1))
    bias = np.zeros((1,))
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=1, padding='same',
        output_padding=None, data_format='channels_last', dilation_rate=1
    ); out = codeflash_output

def test_large_scale_stride_and_dilation():
    # Test 19: Large stride and dilation, moderate size
    x = np.ones((2,4,4,4,2))
    kernel = np.ones((2,2,2,4,2))
    bias = np.zeros((4,))
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=2, padding='same',
        output_padding=None, data_format='channels_last', dilation_rate=2
    ); out = codeflash_output

def test_large_scale_non_cubic():
    # Test 20: Large non-cubic input/output
    x = np.ones((1,5,7,9,3))
    kernel = np.ones((2,3,4,5,3))
    bias = np.ones((5,))
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=(2,2,2), padding='same',
        output_padding=None, data_format='channels_last', dilation_rate=(1,2,1)
    ); out = codeflash_output

def test_large_scale_channels_first():
    # Test 21: Large, channels_first
    x = np.ones((4,3,5,5,5))
    kernel = np.ones((3,3,3,6,3))
    bias = np.zeros((6,))
    codeflash_output = np_conv3d_transpose(
        x, kernel, bias, strides=1, padding='same',
        output_padding=None, data_format='channels_first', dilation_rate=1
    ); out = codeflash_output
# 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-np_conv3d_transpose-maxfz0kh and push.