⚡️ Speed up method ViTSTR.compute_loss by 6%

[codeflash-ai]

📄 6% (0.06x) speedup for ViTSTR.compute_loss in doctr/models/recognition/vitstr/pytorch.py

⏱️ Runtime : 4.25 milliseconds → 4.03 milliseconds (best of 183 runs)

📝 Explanation and details

The optimized code achieves a 5% speedup through three key improvements:

What optimizations were applied:

1. Avoided input mutation: Created seq_len_ instead of modifying the input seq_len tensor in-place, preventing potential memory allocation overhead from tensor mutation
2. More efficient masking: Replaced cce[mask_2d] = 0 with cce.masked_fill_(mask_2d, 0), which uses PyTorch's optimized in-place masking operation
3. Optimized tensor broadcasting: Split the mask creation into row_range = torch.arange(...) and mask_2d = row_range.unsqueeze(0) >= seq_len_.unsqueeze(1) to avoid repeated tensor indexing operations

Why these optimizations work:

• Input mutation avoidance prevents PyTorch from creating defensive tensor copies when the input might be used elsewhere
• masked_fill_ operation is a specialized PyTorch kernel that's faster than general tensor assignment for zeroing masked elements
• Explicit broadcasting reduces the overhead of PyTorch's automatic broadcasting by creating the range tensor once and reusing it

Performance characteristics:
The optimizations show consistent 6-15% improvements across varied sequence lengths and batch sizes, with particularly strong gains on:

• Small batches with varied sequence lengths (8-15% faster)
• Edge cases like zero-length sequences (7-8% faster)
• Large batches still benefit (2-3% faster), showing the optimizations scale well

The changes preserve all original behavior and error handling while delivering measurable performance gains across the full range of typical ViTSTR loss computation scenarios.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from typing import Any

# imports
import pytest  # used for our unit tests
import torch
from doctr.models.recognition.vitstr.pytorch import ViTSTR
from torch import nn
from torch.nn import functional as F

# function to test
# Copyright (C) 2021-2025, Mindee.

# This program is licensed under the Apache License 2.0.
# See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.


class ViTSTRPostProcessor:
    def __init__(self, vocab: str):
        self.vocab = vocab

class _ViTSTR(nn.Module):
    pass
from doctr.models.recognition.vitstr.pytorch import ViTSTR

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

def test_loss_single_perfect_match():
    # Single batch, short sequence, perfect prediction
    batch_size = 1
    seq_len = 4
    num_classes = 6
    # ground truth: [SOS, 1, 2, 3, EOS]
    gt = torch.tensor([[0, 1, 2, 3, 5]])  # EOS index = 5
    seq_len_tensor = torch.tensor([4])  # 4 chars (excluding SOS)
    # logits: for each position, highest at gt label
    logits = torch.full((batch_size, seq_len, num_classes), -10.0)
    for t in range(seq_len):
        logits[0, t, gt[0, t+1]] = 10.0  # position t, label gt[t+1]
    # Compute loss
    codeflash_output = ViTSTR.compute_loss(logits, gt, seq_len_tensor); loss = codeflash_output # 107μs -> 100μs (6.36% faster)

def test_loss_batch_basic():
    # Batch of 2, short sequence, random logits
    batch_size = 2
    seq_len = 3
    num_classes = 5
    gt = torch.tensor([
        [0, 1, 2, 4],  # EOS = 4
        [0, 2, 3, 4]
    ])
    seq_len_tensor = torch.tensor([3, 3])
    logits = torch.randn(batch_size, seq_len, num_classes)
    codeflash_output = ViTSTR.compute_loss(logits, gt, seq_len_tensor); loss = codeflash_output # 119μs -> 111μs (7.40% faster)

def test_loss_varied_seq_lengths():
    # Batch with varied sequence lengths
    batch_size = 3
    seq_len = 5
    num_classes = 7
    gt = torch.tensor([
        [0, 1, 2, 3, 6, 0],  # EOS = 6, seq_len = 4
        [0, 2, 3, 6, 0, 0],  # EOS = 6, seq_len = 3
        [0, 3, 6, 0, 0, 0],  # EOS = 6, seq_len = 2
    ])
    seq_len_tensor = torch.tensor([4, 3, 2])
    logits = torch.randn(batch_size, seq_len, num_classes)
    codeflash_output = ViTSTR.compute_loss(logits, gt, seq_len_tensor); loss = codeflash_output # 124μs -> 114μs (8.89% faster)

def test_loss_reduction_mean():
    # Test that the reduction is mean over batch
    batch_size = 2
    seq_len = 3
    num_classes = 5
    gt = torch.tensor([
        [0, 1, 2, 4],  # EOS = 4
        [0, 2, 3, 4]
    ])
    seq_len_tensor = torch.tensor([3, 3])
    logits = torch.zeros(batch_size, seq_len, num_classes)
    # Set logits so that each batch has different loss
    logits[0] = 2.0  # uniform, low loss
    logits[1] = -2.0 # uniform, higher loss
    codeflash_output = ViTSTR.compute_loss(logits, gt, seq_len_tensor); loss = codeflash_output # 109μs -> 103μs (6.44% faster)
    # Should be between loss for batch 0 and batch 1
    # Compute per-batch losses
    codeflash_output = ViTSTR.compute_loss(logits[0:1], gt[0:1], seq_len_tensor[0:1]); l0 = codeflash_output # 52.6μs -> 46.1μs (14.1% faster)
    codeflash_output = ViTSTR.compute_loss(logits[1:2], gt[1:2], seq_len_tensor[1:2]); l1 = codeflash_output # 40.9μs -> 35.5μs (15.1% faster)

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

def test_loss_eos_masking():
    # EOS token should mask loss after its position
    batch_size = 1
    seq_len = 5
    num_classes = 8
    gt = torch.tensor([[0, 1, 2, 7, 0, 0]])  # EOS=7 at pos 3, seq_len=3
    seq_len_tensor = torch.tensor([3])
    # logits: all random except after EOS, set to high value at wrong class
    logits = torch.randn(batch_size, seq_len, num_classes)
    # After EOS, set logits to have high value at wrong class
    logits[0, 3:, :] = -10.0
    logits[0, 3:, 0] = 10.0
    # Loss should not change if we change logits after EOS
    codeflash_output = ViTSTR.compute_loss(logits, gt, seq_len_tensor); loss1 = codeflash_output # 92.2μs -> 85.6μs (7.80% faster)
    logits[0, 3:, :] = 10.0
    codeflash_output = ViTSTR.compute_loss(logits, gt, seq_len_tensor); loss2 = codeflash_output # 46.9μs -> 41.1μs (13.9% faster)

def test_loss_all_eos_early():
    # All sequences have EOS at position 1 (minimum length)
    batch_size = 3
    seq_len = 5
    num_classes = 6
    gt = torch.tensor([
        [0, 5, 0, 0, 0, 0],  # EOS=5 at pos 1
        [0, 5, 0, 0, 0, 0],
        [0, 5, 0, 0, 0, 0],
    ])
    seq_len_tensor = torch.tensor([1, 1, 1])
    logits = torch.randn(batch_size, seq_len, num_classes)
    codeflash_output = ViTSTR.compute_loss(logits, gt, seq_len_tensor); loss = codeflash_output # 110μs -> 103μs (6.55% faster)

def test_loss_zero_length_sequence():
    # Sequence with length zero (only SOS, EOS at pos 1)
    batch_size = 2
    seq_len = 3
    num_classes = 4
    gt = torch.tensor([
        [0, 3, 0, 0],  # EOS=3 at pos 1
        [0, 3, 0, 0],
    ])
    seq_len_tensor = torch.tensor([1, 1])  # zero-length between SOS and EOS
    logits = torch.randn(batch_size, seq_len, num_classes)
    codeflash_output = ViTSTR.compute_loss(logits, gt, seq_len_tensor); loss = codeflash_output # 111μs -> 104μs (7.35% faster)

def test_loss_invalid_gt_shape():
    # gt tensor shape mismatch
    batch_size = 2
    seq_len = 4
    num_classes = 5
    logits = torch.randn(batch_size, seq_len, num_classes)
    gt = torch.zeros(batch_size, seq_len, dtype=torch.long)  # Should be seq_len+1
    seq_len_tensor = torch.tensor([4, 4])
    with pytest.raises(RuntimeError):
        ViTSTR.compute_loss(logits, gt, seq_len_tensor) # 99.6μs -> 99.8μs (0.120% slower)

def test_loss_invalid_logits_shape():
    # logits tensor shape mismatch
    batch_size = 2
    seq_len = 4
    num_classes = 5
    logits = torch.randn(batch_size, seq_len+1, num_classes)  # Should be seq_len
    gt = torch.zeros(batch_size, seq_len+1, dtype=torch.long)
    seq_len_tensor = torch.tensor([4, 4])
    with pytest.raises(RuntimeError):
        ViTSTR.compute_loss(logits, gt, seq_len_tensor) # 88.5μs -> 90.0μs (1.66% slower)


def test_loss_negative_seq_len():
    # Negative sequence length should cause masking of all positions
    batch_size = 1
    seq_len = 3
    num_classes = 4
    logits = torch.randn(batch_size, seq_len, num_classes)
    gt = torch.zeros(batch_size, seq_len+1, dtype=torch.long)
    seq_len_tensor = torch.tensor([-1])
    # Should not error, but loss should be zero (all masked)
    codeflash_output = ViTSTR.compute_loss(logits, gt, seq_len_tensor); loss = codeflash_output # 135μs -> 126μs (7.74% faster)

def test_loss_extreme_logits():
    # Extremely large logits should produce near-zero loss for correct class
    batch_size = 1
    seq_len = 2
    num_classes = 3
    gt = torch.tensor([[0, 1, 2]])
    seq_len_tensor = torch.tensor([2])
    logits = torch.full((batch_size, seq_len, num_classes), -1000.0)
    for t in range(seq_len):
        logits[0, t, gt[0, t+1]] = 1000.0
    codeflash_output = ViTSTR.compute_loss(logits, gt, seq_len_tensor); loss = codeflash_output # 105μs -> 99.1μs (6.68% faster)

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

def test_loss_large_batch():
    # Large batch size to test scalability
    batch_size = 512
    seq_len = 8
    num_classes = 10
    gt = torch.zeros(batch_size, seq_len+1, dtype=torch.long)
    for i in range(batch_size):
        gt[i, 1:] = torch.randint(0, num_classes-1, (seq_len,))
    seq_len_tensor = torch.full((batch_size,), seq_len)
    logits = torch.randn(batch_size, seq_len, num_classes)
    codeflash_output = ViTSTR.compute_loss(logits, gt, seq_len_tensor); loss = codeflash_output # 336μs -> 326μs (3.13% faster)

def test_loss_long_sequences():
    # Long sequence length (but <1000 elements)
    batch_size = 2
    seq_len = 128
    num_classes = 20
    gt = torch.zeros(batch_size, seq_len+1, dtype=torch.long)
    for i in range(batch_size):
        gt[i, 1:] = torch.randint(0, num_classes-1, (seq_len,))
    seq_len_tensor = torch.full((batch_size,), seq_len)
    logits = torch.randn(batch_size, seq_len, num_classes)
    codeflash_output = ViTSTR.compute_loss(logits, gt, seq_len_tensor); loss = codeflash_output # 119μs -> 111μs (7.39% faster)

def test_loss_max_size_tensors():
    # Test near the allowed memory (not exceeding 100MB)
    batch_size = 16
    seq_len = 256
    num_classes = 24
    gt = torch.zeros(batch_size, seq_len+1, dtype=torch.long)
    for i in range(batch_size):
        gt[i, 1:] = torch.randint(0, num_classes-1, (seq_len,))
    seq_len_tensor = torch.full((batch_size,), seq_len)
    logits = torch.randn(batch_size, seq_len, num_classes)
    codeflash_output = ViTSTR.compute_loss(logits, gt, seq_len_tensor); loss = codeflash_output # 252μs -> 242μs (4.06% faster)

def test_loss_random_large_varied_lengths():
    # Large batch, varied sequence lengths
    batch_size = 256
    seq_len = 32
    num_classes = 12
    gt = torch.zeros(batch_size, seq_len+1, dtype=torch.long)
    seq_len_tensor = torch.randint(1, seq_len+1, (batch_size,))
    for i in range(batch_size):
        gt[i, 1:seq_len_tensor[i]+1] = torch.randint(0, num_classes-1, (seq_len_tensor[i],))
        gt[i, seq_len_tensor[i]+1:] = 0  # pad
    logits = torch.randn(batch_size, seq_len, num_classes)
    codeflash_output = ViTSTR.compute_loss(logits, gt, seq_len_tensor); loss = codeflash_output # 277μs -> 269μs (2.83% faster)

def test_loss_gradients_large():
    # Check that loss is differentiable and gradients can be computed for large batch
    batch_size = 64
    seq_len = 16
    num_classes = 8
    gt = torch.zeros(batch_size, seq_len+1, dtype=torch.long)
    for i in range(batch_size):
        gt[i, 1:] = torch.randint(0, num_classes-1, (seq_len,))
    seq_len_tensor = torch.full((batch_size,), seq_len)
    logits = torch.randn(batch_size, seq_len, num_classes, requires_grad=True)
    codeflash_output = ViTSTR.compute_loss(logits, gt, seq_len_tensor); loss = codeflash_output # 159μs -> 149μs (6.50% faster)
    loss.backward()
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from typing import Any

# imports
import pytest  # used for our unit tests
import torch
from doctr.models.recognition.vitstr.pytorch import ViTSTR
from torch import nn
from torch.nn import functional as F

# function to test
# Copyright (C) 2021-2025, Mindee.

# This program is licensed under the Apache License 2.0.
# See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.



class ViTSTRPostProcessor:
    def __init__(self, vocab: str):
        self.vocab = vocab

class _ViTSTR:
    pass
from doctr.models.recognition.vitstr.pytorch import ViTSTR

# unit tests

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

def test_loss_perfect_prediction():
    # Perfect prediction: logits are always highest for the correct class
    batch_size, seq_len, num_classes = 2, 4, 5
    # gt: shape (batch, seq_len+2) (with <sos> and <eos>)
    gt = torch.tensor([
        [0, 1, 2, 3, 4, 5],  # <sos>, 1,2,3,4,5
        [0, 2, 3, 1, 4, 5]
    ])
    seq_lengths = torch.tensor([4, 4])  # length of word (excluding <sos> and <eos>)
    # model_output: shape (batch, seq_len+1, num_classes+1)
    # Only the correct class has high logits
    logits = torch.full((batch_size, seq_len+1, num_classes+1), -10.0)
    for b in range(batch_size):
        for t in range(seq_len+1):
            target_class = gt[b, t+1]
            logits[b, t, target_class] = 10.0
    # Loss should be near zero
    codeflash_output = ViTSTR.compute_loss(logits, gt, seq_lengths); loss = codeflash_output # 117μs -> 109μs (7.22% faster)

def test_loss_random_prediction():
    # Random logits, loss should be positive
    batch_size, seq_len, num_classes = 2, 4, 5
    gt = torch.tensor([
        [0, 1, 2, 3, 4, 5],
        [0, 2, 3, 1, 4, 5]
    ])
    seq_lengths = torch.tensor([4, 4])
    logits = torch.randn(batch_size, seq_len+1, num_classes+1)
    codeflash_output = ViTSTR.compute_loss(logits, gt, seq_lengths); loss = codeflash_output # 121μs -> 113μs (7.00% faster)

def test_loss_batch_size_one():
    # Single batch
    batch_size, seq_len, num_classes = 1, 3, 4
    gt = torch.tensor([[0, 1, 2, 3, 4]])
    seq_lengths = torch.tensor([3])
    logits = torch.full((batch_size, seq_len+1, num_classes+1), -5.0)
    for t in range(seq_len+1):
        target_class = gt[0, t+1]
        logits[0, t, target_class] = 5.0
    codeflash_output = ViTSTR.compute_loss(logits, gt, seq_lengths); loss = codeflash_output # 102μs -> 97.2μs (5.86% faster)

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

def test_loss_zero_length_sequence():
    # Sequence length zero (only <sos> and <eos>)
    batch_size, seq_len, num_classes = 2, 0, 3
    gt = torch.tensor([
        [0, 1],  # <sos>, <eos>
        [0, 2]
    ])
    seq_lengths = torch.tensor([0, 0])
    logits = torch.zeros(batch_size, seq_len+1, num_classes+1)
    # Should not crash, should return 0 (no steps to predict)
    codeflash_output = ViTSTR.compute_loss(logits, gt, seq_lengths); loss = codeflash_output # 116μs -> 107μs (7.81% faster)

def test_loss_max_length_sequence():
    # Sequence with max length allowed
    batch_size, seq_len, num_classes = 1, 32, 10
    gt = torch.cat([torch.tensor([[0]]), torch.randint(1, num_classes+1, (1, seq_len+1))], dim=1)
    seq_lengths = torch.tensor([32])
    logits = torch.randn(batch_size, seq_len+1, num_classes+1)
    codeflash_output = ViTSTR.compute_loss(logits, gt, seq_lengths); loss = codeflash_output # 109μs -> 100μs (8.53% faster)

def test_loss_eos_masking():
    # Ensure loss is masked after EOS
    batch_size, seq_len, num_classes = 1, 5, 6
    # EOS at position 3 (so seq_len = 3)
    gt = torch.tensor([[0, 1, 2, 6, 0, 0, 0]])  # <sos>, 1,2,<eos>,pad,pad,pad
    seq_lengths = torch.tensor([3])
    logits = torch.full((batch_size, seq_len+1, num_classes+1), -10.0)
    # Only first 3 positions have correct class, rest are padding
    for t in range(seq_len+1):
        target_class = gt[0, t+1]
        logits[0, t, target_class] = 10.0
    # Loss should be near zero, and padding positions should not affect loss
    codeflash_output = ViTSTR.compute_loss(logits, gt, seq_lengths); loss = codeflash_output # 103μs -> 97.3μs (6.74% faster)

def test_loss_incorrect_shape_raises():
    # model_output shape mismatch
    batch_size, seq_len, num_classes = 2, 4, 5
    gt = torch.zeros(batch_size, seq_len+2, dtype=torch.long)
    seq_lengths = torch.tensor([4, 4])
    logits = torch.zeros(batch_size, seq_len, num_classes+1)  # should be seq_len+1
    with pytest.raises(RuntimeError):
        ViTSTR.compute_loss(logits, gt, seq_lengths) # 100μs -> 100μs (0.561% slower)

def test_loss_device_consistency():
    # Test with CUDA if available
    if torch.cuda.is_available():
        batch_size, seq_len, num_classes = 2, 4, 5
        gt = torch.tensor([
            [0, 1, 2, 3, 4, 5],
            [0, 2, 3, 1, 4, 5]
        ], device="cuda")
        seq_lengths = torch.tensor([4, 4], device="cuda")
        logits = torch.full((batch_size, seq_len+1, num_classes+1), -10.0, device="cuda")
        for b in range(batch_size):
            for t in range(seq_len+1):
                target_class = gt[b, t+1]
                logits[b, t, target_class] = 10.0
        codeflash_output = ViTSTR.compute_loss(logits, gt, seq_lengths); loss = codeflash_output

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

def test_loss_large_batch():
    # Large batch size
    batch_size, seq_len, num_classes = 512, 8, 10
    gt = torch.randint(0, num_classes+1, (batch_size, seq_len+2))
    seq_lengths = torch.full((batch_size,), seq_len)
    logits = torch.randn(batch_size, seq_len+1, num_classes+1)
    codeflash_output = ViTSTR.compute_loss(logits, gt, seq_lengths); loss = codeflash_output # 395μs -> 386μs (2.21% faster)

def test_loss_large_sequence():
    # Large sequence length
    batch_size, seq_len, num_classes = 2, 999, 8
    gt = torch.randint(0, num_classes+1, (batch_size, seq_len+2))
    seq_lengths = torch.full((batch_size,), seq_len)
    logits = torch.randn(batch_size, seq_len+1, num_classes+1)
    codeflash_output = ViTSTR.compute_loss(logits, gt, seq_lengths); loss = codeflash_output # 163μs -> 154μs (6.01% faster)

def test_loss_large_num_classes():
    # Large number of classes
    batch_size, seq_len, num_classes = 2, 8, 999
    gt = torch.randint(0, num_classes+1, (batch_size, seq_len+2))
    seq_lengths = torch.full((batch_size,), seq_len)
    logits = torch.randn(batch_size, seq_len+1, num_classes+1)
    codeflash_output = ViTSTR.compute_loss(logits, gt, seq_lengths); loss = codeflash_output # 200μs -> 189μs (5.91% faster)

def test_loss_varied_seq_lengths():
    # Varied sequence lengths in batch
    batch_size, seq_len, num_classes = 4, 6, 7
    gt = torch.zeros(batch_size, seq_len+2, dtype=torch.long)
    seq_lengths = torch.tensor([1, 3, 5, 6])
    for b in range(batch_size):
        # Fill up to seq_lengths[b] with random classes, then EOS
        for t in range(seq_lengths[b]):
            gt[b, t+1] = torch.randint(1, num_classes+1, (1,)).item()
        gt[b, seq_lengths[b]+1] = num_classes  # EOS token
    logits = torch.randn(batch_size, seq_len+1, num_classes+1)
    codeflash_output = ViTSTR.compute_loss(logits, gt, seq_lengths); loss = codeflash_output # 111μs -> 102μs (7.91% faster)

def test_loss_extreme_logits():
    # Extreme logits: all logits very large/small
    batch_size, seq_len, num_classes = 2, 5, 6
    gt = torch.randint(0, num_classes+1, (batch_size, seq_len+2))
    seq_lengths = torch.full((batch_size,), seq_len)
    logits = torch.full((batch_size, seq_len+1, num_classes+1), 1e6)
    codeflash_output = ViTSTR.compute_loss(logits, gt, seq_lengths); loss = codeflash_output # 118μs -> 108μs (9.19% 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-ViTSTR.compute_loss-mg7j7c4h and push.