⚡️ Speed up method ClapTextEmbeddings.create_position_ids_from_input_ids by 22%

[codeflash-ai]

📄 22% (0.22x) speedup for ClapTextEmbeddings.create_position_ids_from_input_ids in src/transformers/models/clap/modeling_clap.py

⏱️ Runtime : 1.14 milliseconds → 933 microseconds (best of 76 runs)

📝 Explanation and details

The optimization refactors the tensor computation pipeline in create_position_ids_from_input_ids to eliminate redundant type conversions and streamline operations:

Key optimizations:

1. Eliminated redundant type conversions: The original code computed mask = input_ids.ne(padding_idx).int() and later used .type_as(mask), creating unnecessary int→float→int conversions. The optimized version keeps mask as a boolean tensor throughout, which PyTorch handles efficiently.

2. Separated operations for clarity and efficiency: Instead of the complex single-line expression (torch.cumsum(mask, dim=1).type_as(mask) + past_key_values_length) * mask, the optimized version breaks this into discrete steps:

   • torch.cumsum(mask, dim=1) operates directly on the boolean mask
   • Conditional addition of past_key_values_length only when needed (avoiding unnecessary addition when it's 0)
   • Final multiplication by the mask

3. Conditional optimization: The if past_key_values_length != 0 check avoids adding zero in the common case where no past key values exist, reducing computation overhead.

Performance impact: The line profiler shows the original bottleneck was the complex single-line operation (48.4% of total time). The optimized version spreads this across simpler operations, with the most expensive being torch.cumsum at 32.8% of total time. This results in a 21% speedup overall.

Test case analysis: The optimization shows consistent improvements across all test scenarios, with particularly strong gains (34-52%) on edge cases like single tokens, empty tensors, and padding-heavy sequences. Large-scale tests show more modest but still meaningful improvements (6-21%), indicating the optimization scales well for production workloads.

The refactoring maintains identical functionality while leveraging PyTorch's efficient boolean operations and eliminating unnecessary type conversions in the tensor computation pipeline.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import torch

from transformers.models.clap.modeling_clap import ClapTextEmbeddings


# unit tests

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


def test_single_sequence_no_padding():
    # Basic: 1D input, no padding
    input_ids = torch.tensor([[1, 2, 3, 4]])
    padding_idx = 0
    expected = torch.tensor([[1, 2, 3, 4]])


def test_single_sequence_with_padding():
    # Basic: 1D input, with padding at the end
    input_ids = torch.tensor([[5, 6, 0, 0]])
    padding_idx = 0
    expected = torch.tensor([[1, 2, 0, 0]])


def test_multi_sequence_mixed_padding():
    # Basic: 2D input, multiple sequences with mixed padding
    input_ids = torch.tensor([[0, 7, 8, 0], [9, 0, 10, 11]])
    padding_idx = 0
    expected = torch.tensor([[0, 1, 2, 0], [1, 0, 1, 2]])


def test_nonzero_padding_idx():
    # Basic: padding_idx not zero
    input_ids = torch.tensor([[3, 3, 4, 5, 3]])
    padding_idx = 3
    expected = torch.tensor([[3, 3, 4, 5, 3]])


def test_all_padding():
    # Basic: All elements are padding
    input_ids = torch.tensor([[0, 0, 0]])
    padding_idx = 0
    expected = torch.tensor([[0, 0, 0]])


def test_no_padding():
    # Basic: No padding in any sequence
    input_ids = torch.tensor([[1, 2, 3], [4, 5, 6]])
    padding_idx = 0
    expected = torch.tensor([[1, 2, 3], [1, 2, 3]])


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


def test_empty_input():
    # Edge: Empty input (0-length sequence)
    input_ids = torch.empty((1, 0), dtype=torch.long)
    padding_idx = 0
    expected = torch.empty((1, 0), dtype=torch.long)
    codeflash_output = ClapTextEmbeddings.create_position_ids_from_input_ids(input_ids, padding_idx)
    out = codeflash_output  # 51.4μs -> 41.7μs (23.2% faster)


def test_all_non_padding_nonzero_padding_idx():
    # Edge: All tokens, no padding, nonzero padding_idx
    input_ids = torch.tensor([[5, 6, 7]])
    padding_idx = 3
    expected = torch.tensor([[4, 5, 6]])


def test_varied_length_sequences():
    # Edge: Batch with varied length (simulate by padding)
    input_ids = torch.tensor([[0, 2, 3, 0, 0], [4, 5, 6, 7, 8]])
    padding_idx = 0
    expected = torch.tensor([[0, 1, 2, 0, 0], [1, 2, 3, 4, 5]])


def test_with_past_key_values_length():
    # Edge: Using past_key_values_length
    input_ids = torch.tensor([[0, 1, 2, 0]])
    padding_idx = 0
    past_key_values_length = 2
    expected = torch.tensor([[0, 3, 4, 0]])


def test_negative_padding_idx():
    # Edge: Negative padding_idx (should work, as torch.ne(-1) is valid)
    input_ids = torch.tensor([[-1, 2, 3]])
    padding_idx = -1
    expected = torch.tensor([[-1, 0, 1]])


def test_high_padding_idx():
    # Edge: Padding index higher than any input id
    input_ids = torch.tensor([[1, 2, 3]])
    padding_idx = 99
    expected = torch.tensor([[100, 101, 102]])


def test_one_token_sequence():
    # Edge: Sequence of length 1, not padding
    input_ids = torch.tensor([[8]])
    padding_idx = 0
    expected = torch.tensor([[1]])


def test_one_token_sequence_padding():
    # Edge: Sequence of length 1, is padding
    input_ids = torch.tensor([[0]])
    padding_idx = 0
    expected = torch.tensor([[0]])


def test_large_batch_and_sequence():
    # Large: batch of 100, sequence length 100
    batch_size = 100
    seq_len = 100
    padding_idx = 0
    # Create input with padding at random positions
    torch.manual_seed(42)
    input_ids = torch.randint(1, 1000, (batch_size, seq_len))
    # Set first 10 tokens of each sequence to padding
    input_ids[:, :10] = padding_idx
    # Expected: positions start at 1 after padding, paddings remain 0
    expected = torch.zeros_like(input_ids)
    for i in range(batch_size):
        count = 0
        for j in range(seq_len):
            if input_ids[i, j] != padding_idx:
                count += 1
                expected[i, j] = count
            else:
                expected[i, j] = 0


def test_large_sequence_with_past_key_values_length():
    # Large: Large sequence, nonzero past_key_values_length
    batch_size = 10
    seq_len = 500
    padding_idx = 0
    past_key_values_length = 7
    input_ids = torch.ones((batch_size, seq_len), dtype=torch.long)
    # Insert padding at random positions
    input_ids[:, ::50] = padding_idx
    # Expected: positions incremented by past_key_values_length, paddings remain 0
    expected = torch.zeros_like(input_ids)
    for i in range(batch_size):
        count = 0
        for j in range(seq_len):
            if input_ids[i, j] != padding_idx:
                count += 1
                expected[i, j] = count + past_key_values_length
            else:
                expected[i, j] = 0
    expected = expected + padding_idx


def test_large_nonzero_padding_idx():
    # Large: Large batch, nonzero padding idx
    batch_size = 50
    seq_len = 20
    padding_idx = 5
    input_ids = torch.randint(6, 100, (batch_size, seq_len))
    # Set last 5 tokens to padding
    input_ids[:, -5:] = padding_idx
    expected = torch.zeros_like(input_ids)
    for i in range(batch_size):
        count = 0
        for j in range(seq_len):
            if input_ids[i, j] != padding_idx:
                count += 1
                expected[i, j] = count + padding_idx
            else:
                expected[i, j] = padding_idx


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

import torch

from transformers.models.clap.modeling_clap import ClapTextEmbeddings


# unit tests

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


def test_basic_single_row_no_padding():
    # Single row, no padding
    input_ids = torch.tensor([[1, 2, 3, 4]])
    padding_idx = 0
    # Positions should be [1,2,3,4], since padding_idx=0, so positions start at 1
    expected = torch.tensor([[1, 2, 3, 4]])
    codeflash_output = ClapTextEmbeddings.create_position_ids_from_input_ids(input_ids, padding_idx)
    output = codeflash_output  # 53.2μs -> 42.9μs (24.1% faster)


def test_basic_single_row_with_padding():
    # Single row, some padding at start and end
    input_ids = torch.tensor([[0, 1, 2, 0]])
    padding_idx = 0
    # Positions: [0,1,2,0]
    expected = torch.tensor([[0, 1, 2, 0]])
    codeflash_output = ClapTextEmbeddings.create_position_ids_from_input_ids(input_ids, padding_idx)
    output = codeflash_output  # 47.0μs -> 37.5μs (25.1% faster)


def test_basic_multiple_rows():
    # Multiple rows, mixed padding
    input_ids = torch.tensor([[0, 1, 2, 0], [1, 0, 2, 3]])
    padding_idx = 0
    # First row: [0,1,2,0]
    # Second row: [1,0,2,3]
    expected = torch.tensor([[0, 1, 2, 0], [1, 0, 2, 3]])
    codeflash_output = ClapTextEmbeddings.create_position_ids_from_input_ids(input_ids, padding_idx)
    output = codeflash_output  # 46.1μs -> 36.5μs (26.3% faster)


def test_basic_nonzero_padding_idx():
    # Padding idx is not zero
    input_ids = torch.tensor([[2, 3, 2, 4]])
    padding_idx = 2
    # Positions: [2,3,2,4] -> mask: [0,1,0,1], position numbers start at 3
    # cumsum(mask): [0,1,1,2], so positions: ([0,1,0,2] + 2) = [2,3,2,4]
    expected = torch.tensor([[2, 3, 2, 4]])
    codeflash_output = ClapTextEmbeddings.create_position_ids_from_input_ids(input_ids, padding_idx)
    output = codeflash_output  # 44.0μs -> 35.1μs (25.6% faster)


def test_basic_past_key_values_length():
    # Test with past_key_values_length
    input_ids = torch.tensor([[0, 1, 2, 0]])
    padding_idx = 0
    past_key_values_length = 2
    # mask: [0,1,1,0]
    # cumsum(mask): [0,1,2,2]
    # incremental_indices: ([0,1,2,0] + 2) * mask = [0,3,4,0]
    # final: [0,3,4,0]
    expected = torch.tensor([[0, 3, 4, 0]])
    codeflash_output = ClapTextEmbeddings.create_position_ids_from_input_ids(
        input_ids, padding_idx, past_key_values_length
    )
    output = codeflash_output  # 41.1μs -> 37.3μs (10.2% faster)


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


def test_edge_all_padding():
    # All tokens are padding
    input_ids = torch.tensor([[0, 0, 0, 0]])
    padding_idx = 0
    # All positions should be 0
    expected = torch.tensor([[0, 0, 0, 0]])
    codeflash_output = ClapTextEmbeddings.create_position_ids_from_input_ids(input_ids, padding_idx)
    output = codeflash_output  # 43.9μs -> 30.9μs (42.1% faster)


def test_edge_empty_tensor():
    # Empty tensor
    input_ids = torch.empty((1, 0), dtype=torch.long)
    padding_idx = 0
    expected = torch.empty((1, 0), dtype=torch.long)
    codeflash_output = ClapTextEmbeddings.create_position_ids_from_input_ids(input_ids, padding_idx)
    output = codeflash_output  # 38.0μs -> 30.0μs (26.8% faster)


def test_edge_single_token():
    # Only one token, not padding
    input_ids = torch.tensor([[5]])
    padding_idx = 0
    expected = torch.tensor([[1]])
    codeflash_output = ClapTextEmbeddings.create_position_ids_from_input_ids(input_ids, padding_idx)
    output = codeflash_output  # 45.9μs -> 34.1μs (34.6% faster)


def test_edge_single_token_padding():
    # Only one token, is padding
    input_ids = torch.tensor([[0]])
    padding_idx = 0
    expected = torch.tensor([[0]])
    codeflash_output = ClapTextEmbeddings.create_position_ids_from_input_ids(input_ids, padding_idx)
    output = codeflash_output  # 45.7μs -> 32.1μs (42.4% faster)


def test_edge_negative_padding_idx():
    # Negative padding_idx
    input_ids = torch.tensor([[-1, 1, 2, -1]])
    padding_idx = -1
    # Positions: [-1,0,1,-1] -> mask: [0,1,1,0], positions start at 0
    # cumsum(mask): [0,1,2,2]
    # ([0,1,2,0] + -1) = [-1,0,1,-1]
    expected = torch.tensor([[-1, 0, 1, -1]])
    codeflash_output = ClapTextEmbeddings.create_position_ids_from_input_ids(input_ids, padding_idx)
    output = codeflash_output  # 41.7μs -> 33.2μs (25.6% faster)


def test_edge_high_padding_idx():
    # Padding idx is high value
    input_ids = torch.tensor([[99, 1, 2, 99]])
    padding_idx = 99
    # mask: [0,1,1,0], positions start at 100
    # cumsum(mask): [0,1,2,2]
    # ([0,1,2,0] + 99) = [99,100,101,99]
    expected = torch.tensor([[99, 100, 101, 99]])
    codeflash_output = ClapTextEmbeddings.create_position_ids_from_input_ids(input_ids, padding_idx)
    output = codeflash_output  # 41.0μs -> 31.7μs (29.0% faster)


def test_edge_past_key_values_length_zero_and_nonzero():
    # Test with past_key_values_length = 0 and nonzero
    input_ids = torch.tensor([[0, 1, 2, 0]])
    padding_idx = 0
    # past_key_values_length=0: [0,1,2,0]
    expected_zero = torch.tensor([[0, 1, 2, 0]])
    codeflash_output = ClapTextEmbeddings.create_position_ids_from_input_ids(input_ids, padding_idx, 0)
    output_zero = codeflash_output  # 40.8μs -> 32.5μs (25.7% faster)

    # past_key_values_length=5: [0,6,7,0]
    expected_five = torch.tensor([[0, 6, 7, 0]])
    codeflash_output = ClapTextEmbeddings.create_position_ids_from_input_ids(input_ids, padding_idx, 5)
    output_five = codeflash_output  # 15.0μs -> 11.2μs (34.2% faster)


def test_edge_dtype_preservation():
    # Input dtype is long, output should be long
    input_ids = torch.tensor([[0, 1, 2, 0]], dtype=torch.long)
    padding_idx = 0
    codeflash_output = ClapTextEmbeddings.create_position_ids_from_input_ids(input_ids, padding_idx)
    output = codeflash_output  # 40.9μs -> 30.7μs (33.1% faster)

    # Input dtype is int32, output should be long (function always returns long)
    input_ids = torch.tensor([[0, 1, 2, 0]], dtype=torch.int32)
    codeflash_output = ClapTextEmbeddings.create_position_ids_from_input_ids(input_ids, padding_idx)
    output = codeflash_output  # 16.7μs -> 12.1μs (37.4% faster)


def test_edge_2d_and_1d_input():
    # 2D input
    input_ids = torch.tensor([[0, 1, 2, 0], [1, 2, 0, 0]])
    padding_idx = 0
    expected = torch.tensor([[0, 1, 2, 0], [1, 2, 0, 0]])
    codeflash_output = ClapTextEmbeddings.create_position_ids_from_input_ids(input_ids, padding_idx)
    output = codeflash_output  # 43.5μs -> 31.6μs (37.7% faster)

    # 1D input
    input_ids_1d = torch.tensor([0, 1, 2, 0])
    expected_1d = torch.tensor([0, 1, 2, 0])
    codeflash_output = ClapTextEmbeddings.create_position_ids_from_input_ids(input_ids_1d.unsqueeze(0), padding_idx)
    output_1d = codeflash_output  # 15.6μs -> 10.2μs (52.7% faster)


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


def test_large_scale_no_padding():
    # Large batch, no padding
    batch_size = 100
    seq_length = 500
    input_ids = torch.arange(1, seq_length + 1).repeat(batch_size, 1)
    padding_idx = 0
    # Each row: [1,2,...,500], positions: [1,2,...,500]
    expected = torch.arange(1, seq_length + 1).repeat(batch_size, 1)
    codeflash_output = ClapTextEmbeddings.create_position_ids_from_input_ids(input_ids, padding_idx)
    output = codeflash_output  # 142μs -> 134μs (6.25% faster)


def test_large_scale_all_padding():
    # Large batch, all padding
    batch_size = 50
    seq_length = 800
    input_ids = torch.full((batch_size, seq_length), 0, dtype=torch.long)
    padding_idx = 0
    expected = torch.zeros((batch_size, seq_length), dtype=torch.long)
    codeflash_output = ClapTextEmbeddings.create_position_ids_from_input_ids(input_ids, padding_idx)
    output = codeflash_output  # 127μs -> 118μs (7.34% faster)


def test_large_scale_mixed_padding():
    # Large batch, mixed padding
    batch_size = 10
    seq_length = 1000
    input_ids = torch.full((batch_size, seq_length), 0, dtype=torch.long)
    # Set every 10th token to non-padding
    for i in range(batch_size):
        input_ids[i, ::10] = torch.arange(1, 101)[: seq_length // 10]
    padding_idx = 0
    # For each row, positions: every 10th token gets increasing position starting at 1, rest are 0
    expected = torch.zeros((batch_size, seq_length), dtype=torch.long)
    for i in range(batch_size):
        expected[i, ::10] = torch.arange(1, seq_length // 10 + 1)
    codeflash_output = ClapTextEmbeddings.create_position_ids_from_input_ids(input_ids, padding_idx)
    output = codeflash_output  # 62.5μs -> 51.5μs (21.4% faster)


def test_large_scale_past_key_values_length():
    # Large batch, no padding, with past_key_values_length
    batch_size = 20
    seq_length = 500
    past_key_values_length = 100
    input_ids = torch.arange(1, seq_length + 1).repeat(batch_size, 1)
    padding_idx = 0
    # Positions: [101,102,...,600]
    expected = torch.arange(1, seq_length + 1).repeat(batch_size, 1) + past_key_values_length
    codeflash_output = ClapTextEmbeddings.create_position_ids_from_input_ids(
        input_ids, padding_idx, past_key_values_length
    )
    output = codeflash_output  # 56.1μs -> 49.9μs (12.4% faster)


def test_large_scale_varied_padding_idx():
    # Large batch, varied padding_idx
    batch_size = 5
    seq_length = 100
    padding_idx = 42
    input_ids = torch.full((batch_size, seq_length), padding_idx, dtype=torch.long)
    # Set every 5th token to non-padding
    for i in range(batch_size):
        input_ids[i, ::5] = torch.arange(1, 21)[: seq_length // 5]
    # Expected: every 5th token: positions start at 43, rest are 42
    expected = torch.full((batch_size, seq_length), padding_idx, dtype=torch.long)
    for i in range(batch_size):
        expected[i, ::5] = torch.arange(1, seq_length // 5 + 1) + padding_idx
    codeflash_output = ClapTextEmbeddings.create_position_ids_from_input_ids(input_ids, padding_idx)
    output = codeflash_output  # 38.2μs -> 27.5μs (38.9% 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-ClapTextEmbeddings.create_position_ids_from_input_ids-misdxw9u and push.