⚡️ Speed up method ClapTextEmbeddings.forward by 6%

[codeflash-ai]

📄 6% (0.06x) speedup for ClapTextEmbeddings.forward in src/transformers/models/clap/modeling_clap.py

⏱️ Runtime : 1.78 milliseconds → 1.67 milliseconds (best of 23 runs)

📝 Explanation and details

The optimized version achieves a 6% speedup through three key optimizations in the forward() method:

1. Optimized Token Type IDs Processing
The original code always expanded self.token_type_ids buffer regardless of whether it was already the correct size. The optimization adds a conditional check to avoid unnecessary expansion:

if self.token_type_ids.shape[0] != bs_pos:
    buffered_token_type_ids = self.token_type_ids.expand(bs_pos, -1)
else:
    buffered_token_type_ids = self.token_type_ids

This reduces redundant tensor operations when the buffer is already correctly sized.

2. Conditional Final Expansion
Similarly, the final expansion to (batch_size, seq_length) is now conditional:

if (batch_size, seq_length) != buffered_token_type_ids.shape:
    token_type_ids = buffered_token_type_ids.expand(batch_size, seq_length)
else:
    token_type_ids = buffered_token_type_ids

This avoids creating new tensors when dimensions already match.

3. In-place Addition for Embeddings
The embedding combination is optimized using .add() method instead of + operator:

embeddings = inputs_embeds
embeddings = embeddings.add(token_type_embeddings)
embeddings = embeddings.add(position_embeddings)

This can be more memory-efficient and potentially faster than creating intermediate tensors.

4. Optimized Position ID Creation
In create_position_ids_from_input_ids(), the optimization eliminates redundant type conversions by working directly with boolean masks and avoiding intermediate .int() conversion, then combining operations more efficiently.

Performance Impact
The line profiler shows the most significant gains in the token type processing logic (lines with expansion operations) and the create_position_ids functions. The test results demonstrate consistent 2-11% improvements across various input sizes and configurations, with larger improvements on smaller inputs where the overhead reduction is more pronounced. This optimization particularly benefits models processing variable-length sequences where conditional expansions can frequently avoid unnecessary work.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

# imports
import pytest
import torch

from transformers.models.clap.modeling_clap import ClapTextEmbeddings


# function to test
# (ClapTextEmbeddings as provided above, assumed present in the same file)


# Helper config class for tests
class DummyConfig:
    def __init__(
        self,
        vocab_size=100,
        hidden_size=32,
        pad_token_id=0,
        type_vocab_size=2,
        layer_norm_eps=1e-5,
        hidden_dropout_prob=0.1,
        max_position_embeddings=128,
    ):
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.pad_token_id = pad_token_id
        self.type_vocab_size = type_vocab_size
        self.layer_norm_eps = layer_norm_eps
        self.hidden_dropout_prob = hidden_dropout_prob
        self.max_position_embeddings = max_position_embeddings


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


def test_forward_basic_input_ids_only():
    """Test basic usage with only input_ids provided (no padding)."""
    config = DummyConfig()
    model = ClapTextEmbeddings(config)
    batch_size, seq_len = 2, 8
    input_ids = torch.randint(1, config.vocab_size, (batch_size, seq_len))
    codeflash_output = model.forward(input_ids=input_ids)
    out = codeflash_output
    # Output should be [batch_size, seq_len, hidden_size]
    assert_tensor_shape_and_dtype(out, (batch_size, seq_len, config.hidden_size), torch.float32)


def test_forward_with_token_type_ids():
    """Test with explicit token_type_ids provided."""
    config = DummyConfig()
    model = ClapTextEmbeddings(config)
    batch_size, seq_len = 3, 12
    input_ids = torch.randint(1, config.vocab_size, (batch_size, seq_len))
    token_type_ids = torch.randint(0, config.type_vocab_size, (batch_size, seq_len))
    codeflash_output = model.forward(input_ids=input_ids, token_type_ids=token_type_ids)
    out = codeflash_output
    assert_tensor_shape_and_dtype(out, (batch_size, seq_len, config.hidden_size), torch.float32)


def test_forward_with_position_ids():
    """Test with explicit position_ids provided."""
    config = DummyConfig()
    model = ClapTextEmbeddings(config)
    batch_size, seq_len = 2, 10
    input_ids = torch.randint(1, config.vocab_size, (batch_size, seq_len))
    position_ids = (
        torch.arange(config.pad_token_id + 1, seq_len + config.pad_token_id + 1)
        .unsqueeze(0)
        .expand(batch_size, seq_len)
    )
    codeflash_output = model.forward(input_ids=input_ids, position_ids=position_ids)
    out = codeflash_output
    assert_tensor_shape_and_dtype(out, (batch_size, seq_len, config.hidden_size), torch.float32)


def test_forward_with_inputs_embeds_only():
    """Test with inputs_embeds provided, no input_ids."""
    config = DummyConfig()
    model = ClapTextEmbeddings(config)
    batch_size, seq_len = 2, 7
    inputs_embeds = torch.randn(batch_size, seq_len, config.hidden_size)
    codeflash_output = model.forward(inputs_embeds=inputs_embeds)
    out = codeflash_output
    assert_tensor_shape_and_dtype(out, (batch_size, seq_len, config.hidden_size), torch.float32)


def test_forward_with_all_arguments():
    """Test with all arguments provided."""
    config = DummyConfig()
    model = ClapTextEmbeddings(config)
    batch_size, seq_len = 2, 6
    input_ids = torch.randint(1, config.vocab_size, (batch_size, seq_len))
    token_type_ids = torch.randint(0, config.type_vocab_size, (batch_size, seq_len))
    position_ids = (
        torch.arange(config.pad_token_id + 1, seq_len + config.pad_token_id + 1)
        .unsqueeze(0)
        .expand(batch_size, seq_len)
    )
    codeflash_output = model.forward(input_ids=input_ids, token_type_ids=token_type_ids, position_ids=position_ids)
    out = codeflash_output
    assert_tensor_shape_and_dtype(out, (batch_size, seq_len, config.hidden_size), torch.float32)


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


def test_forward_empty_sequence():
    """Test with sequence length zero (empty input)."""
    config = DummyConfig()
    model = ClapTextEmbeddings(config)
    batch_size, seq_len = 2, 0
    input_ids = torch.empty(batch_size, seq_len, dtype=torch.long)
    codeflash_output = model.forward(input_ids=input_ids)
    out = codeflash_output
    assert_tensor_shape_and_dtype(out, (batch_size, seq_len, config.hidden_size), torch.float32)


def test_forward_all_padding():
    """Test with input_ids all set to pad_token_id."""
    config = DummyConfig()
    model = ClapTextEmbeddings(config)
    batch_size, seq_len = 2, 5
    input_ids = torch.full((batch_size, seq_len), config.pad_token_id, dtype=torch.long)
    codeflash_output = model.forward(input_ids=input_ids)
    out = codeflash_output
    assert_tensor_shape_and_dtype(out, (batch_size, seq_len, config.hidden_size), torch.float32)
    # All positions should get same position embedding (since all are pad)
    pos_emb = model.position_embeddings(torch.full((1, seq_len), config.pad_token_id, dtype=torch.long))


def test_forward_mixed_padding():
    """Test with a mix of padding and non-padding tokens."""
    config = DummyConfig()
    model = ClapTextEmbeddings(config)
    batch_size, seq_len = 2, 6
    input_ids = torch.tensor(
        [
            [config.pad_token_id, 1, config.pad_token_id, 2, 3, config.pad_token_id],
            [4, config.pad_token_id, 5, config.pad_token_id, config.pad_token_id, 6],
        ]
    )
    codeflash_output = model.forward(input_ids=input_ids)
    out = codeflash_output
    assert_tensor_shape_and_dtype(out, (batch_size, seq_len, config.hidden_size), torch.float32)


def test_forward_extreme_hidden_size():
    """Test with a very large hidden size (within memory constraint)."""
    config = DummyConfig(hidden_size=1024)
    model = ClapTextEmbeddings(config)
    batch_size, seq_len = 2, 10
    input_ids = torch.randint(1, config.vocab_size, (batch_size, seq_len))
    codeflash_output = model.forward(input_ids=input_ids)
    out = codeflash_output
    assert_tensor_shape_and_dtype(out, (batch_size, seq_len, config.hidden_size), torch.float32)


def test_forward_past_key_values_length():
    """Test with nonzero past_key_values_length."""
    config = DummyConfig()
    model = ClapTextEmbeddings(config)
    batch_size, seq_len = 2, 5
    input_ids = torch.randint(1, config.vocab_size, (batch_size, seq_len))
    codeflash_output = model.forward(input_ids=input_ids, past_key_values_length=0)
    out_0 = codeflash_output  # 144μs -> 135μs (6.97% faster)
    codeflash_output = model.forward(input_ids=input_ids, past_key_values_length=2)
    out_2 = codeflash_output  # 72.4μs -> 68.6μs (5.42% faster)


def test_forward_invalid_token_type_ids_shape():
    """Test with token_type_ids of invalid shape (should raise RuntimeError)."""
    config = DummyConfig()
    model = ClapTextEmbeddings(config)
    batch_size, seq_len = 2, 6
    input_ids = torch.randint(1, config.vocab_size, (batch_size, seq_len))
    # Wrong shape: should be (batch_size, seq_len)
    token_type_ids = torch.randint(0, config.type_vocab_size, (batch_size, seq_len + 1))
    with pytest.raises(RuntimeError):
        model.forward(input_ids=input_ids, token_type_ids=token_type_ids)  # 130μs -> 128μs (2.04% faster)


def test_forward_invalid_position_ids_shape():
    """Test with position_ids of invalid shape (should raise RuntimeError)."""
    config = DummyConfig()
    model = ClapTextEmbeddings(config)
    batch_size, seq_len = 2, 6
    input_ids = torch.randint(1, config.vocab_size, (batch_size, seq_len))
    # Wrong shape: should be (batch_size, seq_len)
    position_ids = torch.arange(config.pad_token_id + 1, seq_len + config.pad_token_id + 1).unsqueeze(0)
    with pytest.raises(RuntimeError):
        model.forward(input_ids=input_ids, position_ids=position_ids.expand(batch_size, seq_len + 1))


def test_forward_inputs_embeds_wrong_shape():
    """Test with inputs_embeds of wrong shape (should raise IndexError)."""
    config = DummyConfig()
    model = ClapTextEmbeddings(config)
    # Should be (batch_size, seq_len, hidden_size)
    inputs_embeds = torch.randn(2, config.hidden_size, 10)
    with pytest.raises(IndexError):
        model.forward(inputs_embeds=inputs_embeds)


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


def test_forward_large_batch_and_seq_len():
    """Test with large batch and sequence length, but <100MB tensor."""
    config = DummyConfig(hidden_size=32, max_position_embeddings=256, vocab_size=1000)
    model = ClapTextEmbeddings(config)
    batch_size, seq_len = 64, 64  # 64*64*32*4 bytes = 512KB
    input_ids = torch.randint(1, config.vocab_size, (batch_size, seq_len))
    codeflash_output = model.forward(input_ids=input_ids)
    out = codeflash_output
    assert_tensor_shape_and_dtype(out, (batch_size, seq_len, config.hidden_size), torch.float32)


def test_forward_large_hidden_size():
    """Test with large hidden size, but <100MB tensor."""
    config = DummyConfig(hidden_size=2048, vocab_size=5000, max_position_embeddings=128)
    model = ClapTextEmbeddings(config)
    batch_size, seq_len = 8, 32  # 8*32*2048*4 bytes = 2MB
    input_ids = torch.randint(1, config.vocab_size, (batch_size, seq_len))
    codeflash_output = model.forward(input_ids=input_ids)
    out = codeflash_output
    assert_tensor_shape_and_dtype(out, (batch_size, seq_len, config.hidden_size), torch.float32)


def test_forward_large_inputs_embeds():
    """Test with large inputs_embeds directly, no input_ids."""
    config = DummyConfig(hidden_size=512, max_position_embeddings=256)
    model = ClapTextEmbeddings(config)
    batch_size, seq_len = 16, 128  # 16*128*512*4 bytes = 4MB
    inputs_embeds = torch.randn(batch_size, seq_len, config.hidden_size)
    codeflash_output = model.forward(inputs_embeds=inputs_embeds)
    out = codeflash_output
    assert_tensor_shape_and_dtype(out, (batch_size, seq_len, config.hidden_size), torch.float32)


def test_forward_large_type_vocab_size():
    """Test with large type_vocab_size."""
    config = DummyConfig(type_vocab_size=64)
    model = ClapTextEmbeddings(config)
    batch_size, seq_len = 4, 32
    input_ids = torch.randint(1, config.vocab_size, (batch_size, seq_len))
    token_type_ids = torch.randint(0, config.type_vocab_size, (batch_size, seq_len))
    codeflash_output = model.forward(input_ids=input_ids, token_type_ids=token_type_ids)
    out = codeflash_output
    assert_tensor_shape_and_dtype(out, (batch_size, seq_len, config.hidden_size), torch.float32)


def test_forward_large_max_position_embeddings():
    """Test with large max_position_embeddings."""
    config = DummyConfig(max_position_embeddings=512)
    model = ClapTextEmbeddings(config)
    batch_size, seq_len = 2, 512
    input_ids = torch.randint(1, config.vocab_size, (batch_size, seq_len))
    codeflash_output = model.forward(input_ids=input_ids)
    out = codeflash_output
    assert_tensor_shape_and_dtype(out, (batch_size, seq_len, config.hidden_size), torch.float32)


# 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
import torch

from transformers.models.clap.modeling_clap import ClapTextEmbeddings


# function to test
# (ClapTextEmbeddings class as provided above)


# Dummy config class for testing
class DummyConfig:
    def __init__(
        self,
        vocab_size=10,
        hidden_size=8,
        pad_token_id=0,
        type_vocab_size=2,
        layer_norm_eps=1e-5,
        hidden_dropout_prob=0.1,
        max_position_embeddings=16,
    ):
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.pad_token_id = pad_token_id
        self.type_vocab_size = type_vocab_size
        self.layer_norm_eps = layer_norm_eps
        self.hidden_dropout_prob = hidden_dropout_prob
        self.max_position_embeddings = max_position_embeddings


# Helper to create model
def make_model(config=None):
    if config is None:
        config = DummyConfig()
    return ClapTextEmbeddings(config)


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


def test_forward_basic_input_ids_only():
    # Test with input_ids only, no padding
    model = make_model()
    input_ids = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8]])
    codeflash_output = model.forward(input_ids=input_ids)
    out = codeflash_output  # 142μs -> 131μs (8.31% faster)


def test_forward_with_token_type_ids():
    # Test with input_ids and token_type_ids
    model = make_model()
    input_ids = torch.tensor([[1, 2, 3], [4, 5, 6]])
    token_type_ids = torch.tensor([[0, 1, 0], [1, 0, 1]])
    codeflash_output = model.forward(input_ids=input_ids, token_type_ids=token_type_ids)
    out = codeflash_output  # 124μs -> 115μs (7.17% faster)


def test_forward_with_position_ids():
    # Test with explicit position_ids
    model = make_model()
    input_ids = torch.tensor([[1, 2, 3]])
    position_ids = torch.tensor([[5, 6, 7]])
    codeflash_output = model.forward(input_ids=input_ids, position_ids=position_ids)
    out = codeflash_output  # 103μs -> 100μs (3.46% faster)


def test_forward_with_inputs_embeds_only():
    # Test with inputs_embeds only (no input_ids)
    model = make_model()
    # Create random embeddings of shape (batch_size, seq_length, hidden_size)
    inputs_embeds = torch.randn(2, 5, model.word_embeddings.embedding_dim)
    codeflash_output = model.forward(inputs_embeds=inputs_embeds)
    out = codeflash_output  # 116μs -> 112μs (2.96% faster)


def test_forward_with_all_inputs():
    # Test with all inputs provided
    model = make_model()
    input_ids = torch.tensor([[1, 2], [3, 4]])
    token_type_ids = torch.tensor([[0, 1], [1, 0]])
    position_ids = torch.tensor([[1, 2], [3, 4]])
    inputs_embeds = model.word_embeddings(input_ids)
    codeflash_output = model.forward(
        input_ids=input_ids,
        token_type_ids=token_type_ids,
        position_ids=position_ids,
        inputs_embeds=inputs_embeds,
        past_key_values_length=2,
    )
    out = codeflash_output  # 71.4μs -> 70.0μs (1.98% faster)


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


def test_forward_empty_input():
    # Test with empty input_ids (zero sequence length)
    model = make_model()
    input_ids = torch.empty((2, 0), dtype=torch.long)
    codeflash_output = model.forward(input_ids=input_ids)
    out = codeflash_output  # 106μs -> 96.9μs (10.0% faster)


def test_forward_all_padding():
    # Test with all tokens as padding
    model = make_model()
    input_ids = torch.full((2, 5), model.padding_idx, dtype=torch.long)
    codeflash_output = model.forward(input_ids=input_ids)
    out = codeflash_output  # 136μs -> 127μs (6.91% faster)


def test_forward_mismatched_token_type_ids_shape():
    # Test with token_type_ids shape mismatch
    model = make_model()
    input_ids = torch.tensor([[1, 2, 3]])
    token_type_ids = torch.tensor([[0, 1]])  # Wrong shape
    with pytest.raises(RuntimeError):
        model.forward(input_ids=input_ids, token_type_ids=token_type_ids)  # 131μs -> 127μs (2.83% faster)


def test_forward_position_ids_out_of_bounds():
    # Test with position_ids exceeding max_position_embeddings
    model = make_model()
    input_ids = torch.tensor([[1, 2, 3]])
    position_ids = torch.tensor([[100, 101, 102]])  # Out of bounds
    with pytest.raises(IndexError):
        model.forward(input_ids=input_ids, position_ids=position_ids)


def test_forward_inputs_embeds_wrong_shape():
    # Test with inputs_embeds of wrong shape
    model = make_model()
    inputs_embeds = torch.randn(2, 5)  # Should be (batch, seq, hidden)
    with pytest.raises(IndexError):
        model.forward(inputs_embeds=inputs_embeds)  # 3.71μs -> 3.63μs (2.29% faster)


def test_forward_past_key_values_length_nonzero():
    # Test that past_key_values_length shifts position ids correctly
    model = make_model()
    input_ids = torch.tensor([[1, 2, 0, 3, 0]])
    codeflash_output = model.forward(input_ids=input_ids, past_key_values_length=0)
    out1 = codeflash_output  # 142μs -> 128μs (10.4% faster)
    codeflash_output = model.forward(input_ids=input_ids, past_key_values_length=2)
    out2 = codeflash_output  # 70.2μs -> 63.2μs (11.1% faster)


def test_forward_inputs_embeds_padding_idx():
    # Test create_position_ids_from_inputs_embeds with nonzero padding_idx
    config = DummyConfig(pad_token_id=2)
    model = make_model(config)
    inputs_embeds = torch.randn(1, 4, config.hidden_size)
    pos_ids = model.create_position_ids_from_inputs_embeds(inputs_embeds, config.pad_token_id)


def test_forward_create_position_ids_from_input_ids_padding():
    # Test create_position_ids_from_input_ids with padding
    config = DummyConfig(pad_token_id=1)
    model = make_model(config)
    input_ids = torch.tensor([[2, 1, 3, 1, 4]])
    pos_ids = model.create_position_ids_from_input_ids(input_ids, config.pad_token_id)


def test_forward_token_type_ids_buffer_expansion():
    # Test token_type_ids buffer expansion for batch sizes > 1
    model = make_model()
    input_ids = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
    codeflash_output = model.forward(input_ids=input_ids)
    out = codeflash_output  # 139μs -> 129μs (7.32% faster)


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


def test_forward_large_batch_and_seq():
    # Test with large batch and sequence length (but < 100MB)
    config = DummyConfig(
        vocab_size=100,
        hidden_size=32,
        pad_token_id=0,
        type_vocab_size=2,
        max_position_embeddings=128,
    )
    model = make_model(config)
    batch_size = 64
    seq_length = 128
    input_ids = torch.randint(1, config.vocab_size, (batch_size, seq_length))
    codeflash_output = model.forward(input_ids=input_ids)
    out = codeflash_output


def test_forward_large_vocab_and_hidden():
    # Test with large vocab and hidden size
    config = DummyConfig(
        vocab_size=512,
        hidden_size=64,
        pad_token_id=0,
        type_vocab_size=4,
        max_position_embeddings=256,
    )
    model = make_model(config)
    batch_size = 8
    seq_length = 256
    input_ids = torch.randint(1, config.vocab_size, (batch_size, seq_length))
    codeflash_output = model.forward(input_ids=input_ids)
    out = codeflash_output


def test_forward_large_inputs_embeds():
    # Test with large inputs_embeds directly
    config = DummyConfig(
        vocab_size=256,
        hidden_size=32,
        pad_token_id=0,
        type_vocab_size=2,
        max_position_embeddings=512,
    )
    model = make_model(config)
    batch_size = 16
    seq_length = 512
    inputs_embeds = torch.randn(batch_size, seq_length, config.hidden_size)
    codeflash_output = model.forward(inputs_embeds=inputs_embeds)
    out = codeflash_output


def test_forward_large_type_vocab():
    # Test with large type_vocab_size
    config = DummyConfig(
        vocab_size=32,
        hidden_size=16,
        pad_token_id=0,
        type_vocab_size=32,
        max_position_embeddings=64,
    )
    model = make_model(config)
    batch_size = 32
    seq_length = 64
    input_ids = torch.randint(1, config.vocab_size, (batch_size, seq_length))
    token_type_ids = torch.randint(0, config.type_vocab_size, (batch_size, seq_length))
    codeflash_output = model.forward(input_ids=input_ids, token_type_ids=token_type_ids)
    out = codeflash_output


def test_forward_large_dropout():
    # Test with high dropout to check stability
    config = DummyConfig(hidden_dropout_prob=0.9)
    model = make_model(config)
    input_ids = torch.tensor([[1, 2, 3, 4]])
    codeflash_output = model.forward(input_ids=input_ids)
    out = codeflash_output  # 141μs -> 129μs (9.51% 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.forward-misdcdyz and push.