⚡️ Speed up method RandomThinPlateSpline.generate_parameters by 21%

[codeflash-ai]

📄 21% (0.21x) speedup for RandomThinPlateSpline.generate_parameters in kornia/augmentation/_2d/geometric/thin_plate_spline.py

⏱️ Runtime : 7.44 milliseconds → 6.14 milliseconds (best of 5 runs)

📝 Explanation and details

The optimized code achieves a 21% speedup by eliminating redundant tensor creation in the hot path.

Key Optimization:
The source control points template (a fixed 5x2 tensor with values [[-1,-1], [-1,1], [1,-1], [1,1], [0,0]]) was previously created from scratch on every call to generate_parameters(). The optimization pre-creates this tensor once during __init__ and stores it as self._src_template, then simply copies it to the target device/dtype on each call.

Why This Is Faster:

• Reduced object creation overhead: torch.tensor() involves parsing Python lists, allocating memory, and initializing data. By doing this once instead of per-call, we eliminate ~17-18% of the function's time (line profiler shows the original torch.tensor() call took 17.2% + 10.8% = 28% total time).
• Simpler operation path: The .to() method on an existing tensor is faster than constructing a new tensor from Python literals.
• Memory efficiency: Only one template tensor exists in memory instead of creating temporary tensors per call.

Performance Characteristics:

• The optimization is most effective for workloads with frequent calls to generate_parameters() - evident from test cases showing 32-74% speedup on repeated calls (e.g., test_generate_parameters_repeatability_same_input shows 42.4% faster on second call).
• Batch size agnostic: The speedup is consistent across different batch sizes since the template is expanded, not the creation overhead.
• Minimal impact on edge cases: Tests with batch_size=0 show slight slowdown (6.2%), but this is negligible compared to typical use cases.

Impact on Workloads:
Since generate_parameters() is called during augmentation pipelines, this optimization directly reduces latency in data preprocessing - particularly valuable in training loops where augmentations are applied per-batch. The 21% speedup translates to faster data loading without any change to augmentation quality or behavior.

✅ Correctness verification report:

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

⚙️ Click to see Existing Unit Tests

🌀 Click to see Generated Regression Tests

import pytest
import torch
from kornia.augmentation._2d.geometric.thin_plate_spline import \
    RandomThinPlateSpline

# =========================
# Basic Test Cases
# =========================

def test_generate_parameters_basic_batch_1():
    # Test with batch size 1, typical shape
    aug = RandomThinPlateSpline(scale=0.2)
    shape = (1, 3, 32, 32)
    codeflash_output = aug.generate_parameters(shape); params = codeflash_output # 157μs -> 143μs (9.60% faster)
    src = params["src"]
    dst = params["dst"]
    # Check src is the fixed control points
    expected_src = torch.tensor(
        [[[-1.0, -1.0], [-1.0, 1.0], [1.0, -1.0], [1.0, 1.0], [0.0, 0.0]]]
    )

def test_generate_parameters_basic_batch_4():
    # Test with batch size 4, typical shape
    aug = RandomThinPlateSpline(scale=0.1)
    shape = (4, 3, 16, 16)
    codeflash_output = aug.generate_parameters(shape); params = codeflash_output # 102μs -> 86.8μs (18.6% faster)
    src = params["src"]
    dst = params["dst"]
    # Each batch's src should be the same
    for i in range(1, 4):
        pass
    # dst should differ from src per batch
    for i in range(4):
        pass

def test_generate_parameters_same_on_batch_true():
    # Test that 'same_on_batch' produces identical dst for all batch elements
    aug = RandomThinPlateSpline(scale=0.1, same_on_batch=True)
    shape = (3, 3, 8, 8)
    codeflash_output = aug.generate_parameters(shape); params = codeflash_output # 84.9μs -> 82.1μs (3.42% faster)
    dst = params["dst"]
    # All dst should be identical
    for i in range(1, 3):
        pass

def test_generate_parameters_same_on_batch_false():
    # Test that 'same_on_batch' = False gives different dst for different batch elements (with high probability)
    aug = RandomThinPlateSpline(scale=0.2, same_on_batch=False)
    shape = (8, 3, 8, 8)
    codeflash_output = aug.generate_parameters(shape); params = codeflash_output # 78.7μs -> 69.7μs (12.9% faster)
    dst = params["dst"]
    # With high probability, at least two dsts are different
    found_diff = False
    for i in range(1, 8):
        if not torch.allclose(dst[0], dst[i], atol=1e-6):
            found_diff = True
            break

def test_generate_parameters_scale_zero():
    # Test that scale=0 produces dst == src
    aug = RandomThinPlateSpline(scale=0.0)
    shape = (2, 3, 8, 8)
    codeflash_output = aug.generate_parameters(shape); params = codeflash_output # 77.3μs -> 68.6μs (12.6% faster)
    src, dst = params["src"], params["dst"]

# =========================
# Edge Test Cases
# =========================

def test_generate_parameters_batch_size_zero():
    # Test with batch size zero
    aug = RandomThinPlateSpline(scale=0.2)
    shape = (0, 3, 8, 8)
    codeflash_output = aug.generate_parameters(shape); params = codeflash_output # 70.2μs -> 74.9μs (6.21% slower)
    src, dst = params["src"], params["dst"]

def test_generate_parameters_non_square_image():
    # Test with non-square image shape
    aug = RandomThinPlateSpline(scale=0.2)
    shape = (2, 3, 64, 32)
    codeflash_output = aug.generate_parameters(shape); params = codeflash_output # 75.3μs -> 64.9μs (16.1% faster)
    src, dst = params["src"], params["dst"]

def test_generate_parameters_large_scale():
    # Test with a large scale value
    scale = 5.0
    aug = RandomThinPlateSpline(scale=scale)
    shape = (1, 3, 8, 8)
    codeflash_output = aug.generate_parameters(shape); params = codeflash_output # 74.4μs -> 70.1μs (6.25% faster)
    src, dst = params["src"], params["dst"]
    # Noise should be within [-scale, scale]
    diff = dst - src

def test_generate_parameters_negative_scale():
    # Test with negative scale (should behave as scale=abs(scale))
    aug = RandomThinPlateSpline(scale=-0.5)
    shape = (2, 3, 8, 8)
    codeflash_output = aug.generate_parameters(shape); params = codeflash_output # 75.1μs -> 67.8μs (10.8% faster)
    src, dst = params["src"], params["dst"]
    diff = dst - src

def test_generate_parameters_dtype_and_device():
    # Test that output is on the specified device and dtype
    device = torch.device("cpu")
    dtype = torch.float64
    aug = RandomThinPlateSpline(scale=0.2)
    aug.set_rng_device_and_dtype(device, dtype)
    shape = (2, 3, 8, 8)
    codeflash_output = aug.generate_parameters(shape); params = codeflash_output # 106μs -> 86.5μs (22.8% faster)
    src, dst = params["src"], params["dst"]

def test_generate_parameters_minimal_shape():
    # Test with minimal valid image shape (1x1)
    aug = RandomThinPlateSpline(scale=0.2)
    shape = (1, 1, 1, 1)
    codeflash_output = aug.generate_parameters(shape); params = codeflash_output # 73.3μs -> 65.1μs (12.6% faster)
    src, dst = params["src"], params["dst"]

def test_generate_parameters_invalid_shape_too_small():
    # Test with shape tuple too short (should raise ValueError)
    aug = RandomThinPlateSpline(scale=0.2)
    with pytest.raises(ValueError):
        aug.generate_parameters((1, 3, 8)) # 5.86μs -> 6.49μs (9.69% slower)

def test_generate_parameters_invalid_shape_non_tuple():
    # Test with non-tuple shape (should raise TypeError)
    aug = RandomThinPlateSpline(scale=0.2)
    with pytest.raises(TypeError):
        aug.generate_parameters([1, 3, 8, 8])  # list, not tuple

# =========================
# Large Scale Test Cases
# =========================

def test_generate_parameters_large_batch():
    # Test with a large batch size (within 100MB constraint)
    batch_size = 512  # 512 x 5 x 2 x 4 bytes = 20KB
    aug = RandomThinPlateSpline(scale=0.2)
    shape = (batch_size, 3, 32, 32)
    codeflash_output = aug.generate_parameters(shape); params = codeflash_output # 111μs -> 103μs (8.03% faster)
    src, dst = params["src"], params["dst"]

def test_generate_parameters_large_image():
    # Test with a large image shape (within 100MB constraint)
    aug = RandomThinPlateSpline(scale=0.2)
    shape = (2, 3, 512, 512)
    codeflash_output = aug.generate_parameters(shape); params = codeflash_output # 74.0μs -> 66.1μs (11.9% faster)
    src, dst = params["src"], params["dst"]

def test_generate_parameters_repeatability_same_on_batch():
    # Test that with same_on_batch=True, repeated calls give different results (randomness preserved)
    aug = RandomThinPlateSpline(scale=0.2, same_on_batch=True)
    shape = (4, 3, 32, 32)
    codeflash_output = aug.generate_parameters(shape); params1 = codeflash_output # 81.6μs -> 73.8μs (10.5% faster)
    codeflash_output = aug.generate_parameters(shape); params2 = codeflash_output # 48.5μs -> 36.8μs (32.0% faster)

def test_generate_parameters_repeatability_same_input():
    # Test that two calls with same input shape produce different dsts (randomness preserved)
    aug = RandomThinPlateSpline(scale=0.2)
    shape = (2, 3, 8, 8)
    codeflash_output = aug.generate_parameters(shape); params1 = codeflash_output # 73.2μs -> 80.0μs (8.49% slower)
    codeflash_output = aug.generate_parameters(shape); params2 = codeflash_output # 54.5μs -> 38.2μs (42.4% faster)

# =========================
# Mutation Testing: Ensure src is not mutated
# =========================

def test_generate_parameters_src_is_constant():
    # src should always be the canonical control points, not mutated
    aug = RandomThinPlateSpline(scale=0.2)
    shape = (3, 3, 8, 8)
    codeflash_output = aug.generate_parameters(shape); params1 = codeflash_output # 72.4μs -> 65.2μs (11.0% faster)
    codeflash_output = aug.generate_parameters(shape); params2 = codeflash_output # 57.7μs -> 33.2μs (73.8% faster)
    expected_src = torch.tensor(
        [[[-1.0, -1.0], [-1.0, 1.0], [1.0, -1.0], [1.0, 1.0], [0.0, 0.0]]]
    ).expand(3, 5, 2)
# 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 kornia.augmentation._2d.geometric.thin_plate_spline import \
    RandomThinPlateSpline
from kornia.constants import SamplePadding

# Test Basic Functionality
def test_generate_parameters_returns_dict():
    """Test that generate_parameters returns a dictionary with correct keys."""
    # Create an instance of RandomThinPlateSpline
    aug = RandomThinPlateSpline(scale=0.2, p=1.0)
    
    # Define input shape (B, C, H, W)
    shape = (4, 3, 32, 32)
    
    # Generate parameters
    codeflash_output = aug.generate_parameters(shape); params = codeflash_output # 76.5μs -> 66.5μs (15.0% faster)

def test_generate_parameters_tensor_shapes():
    """Test that generated tensors have correct shapes."""
    # Create an instance with scale=0.2
    aug = RandomThinPlateSpline(scale=0.2, p=1.0)
    
    # Test with batch size 4
    shape = (4, 3, 64, 64)
    codeflash_output = aug.generate_parameters(shape); params = codeflash_output # 74.4μs -> 67.7μs (9.98% faster)

def test_generate_parameters_batch_size_one():
    """Test with batch size of 1."""
    # Create instance
    aug = RandomThinPlateSpline(scale=0.1, p=1.0)
    
    # Test with batch size 1
    shape = (1, 3, 32, 32)
    codeflash_output = aug.generate_parameters(shape); params = codeflash_output # 74.3μs -> 64.9μs (14.6% faster)

def test_generate_parameters_src_values():
    """Test that source control points have expected fixed values."""
    # Create instance
    aug = RandomThinPlateSpline(scale=0.2, p=1.0)
    
    # Generate parameters
    shape = (2, 3, 32, 32)
    codeflash_output = aug.generate_parameters(shape); params = codeflash_output # 91.7μs -> 64.2μs (42.9% faster)
    
    # Expected source points (5 control points in normalized coordinates)
    expected_src = torch.tensor(
        [[[-1.0, -1.0], [-1.0, 1.0], [1.0, -1.0], [1.0, 1.0], [0.0, 0.0]]]
    )
    
    # Verify that all batch items have the same source points
    for i in range(2):
        pass

def test_generate_parameters_dst_is_src_plus_noise():
    """Test that destination points are source points plus noise."""
    # Create instance with known scale
    aug = RandomThinPlateSpline(scale=0.3, p=1.0)
    
    # Generate parameters
    shape = (3, 3, 32, 32)
    codeflash_output = aug.generate_parameters(shape); params = codeflash_output # 73.4μs -> 65.5μs (12.2% faster)
    
    # Calculate the difference (noise)
    noise = params["dst"] - params["src"]

def test_generate_parameters_same_on_batch_true():
    """Test that same_on_batch=True produces identical noise across batch."""
    # Create instance with same_on_batch=True
    aug = RandomThinPlateSpline(scale=0.2, same_on_batch=True, p=1.0)
    
    # Generate parameters with batch size > 1
    shape = (5, 3, 32, 32)
    codeflash_output = aug.generate_parameters(shape); params = codeflash_output # 88.4μs -> 72.6μs (21.7% faster)
    
    # Calculate noise for each batch item
    noise = params["dst"] - params["src"]
    
    # Verify that all batch items have identical noise
    for i in range(1, 5):
        pass

def test_generate_parameters_same_on_batch_false():
    """Test that same_on_batch=False produces different noise across batch."""
    # Create instance with same_on_batch=False
    aug = RandomThinPlateSpline(scale=0.2, same_on_batch=False, p=1.0)
    
    # Generate parameters with batch size > 1
    shape = (5, 3, 32, 32)
    codeflash_output = aug.generate_parameters(shape); params = codeflash_output # 91.1μs -> 67.3μs (35.4% faster)
    
    # Calculate noise for each batch item
    noise = params["dst"] - params["src"]
    
    # With high probability, at least some batch items should have different noise
    # We check if any pair of batch items has different noise
    has_different_noise = False
    for i in range(1, 5):
        if not torch.allclose(noise[0], noise[i], atol=1e-6):
            has_different_noise = True
            break

# Test Edge Cases
def test_generate_parameters_zero_scale():
    """Test with scale=0 (no noise)."""
    # Create instance with scale=0
    aug = RandomThinPlateSpline(scale=0.0, p=1.0)
    
    # Generate parameters
    shape = (3, 3, 32, 32)
    codeflash_output = aug.generate_parameters(shape); params = codeflash_output # 73.3μs -> 64.9μs (12.9% faster)

def test_generate_parameters_large_scale():
    """Test with large scale value."""
    # Create instance with large scale
    aug = RandomThinPlateSpline(scale=5.0, p=1.0)
    
    # Generate parameters
    shape = (2, 3, 32, 32)
    codeflash_output = aug.generate_parameters(shape); params = codeflash_output # 73.1μs -> 64.9μs (12.7% faster)
    
    # Calculate noise
    noise = params["dst"] - params["src"]

def test_generate_parameters_different_image_sizes():
    """Test with various image dimensions."""
    # Create instance
    aug = RandomThinPlateSpline(scale=0.2, p=1.0)
    
    # Test with different image sizes
    sizes = [(2, 3, 16, 16), (3, 1, 64, 64), (1, 3, 128, 256), (4, 3, 224, 224)]
    
    for shape in sizes:
        codeflash_output = aug.generate_parameters(shape); params = codeflash_output # 200μs -> 158μs (26.1% faster)

def test_generate_parameters_dtype_preservation():
    """Test that output tensors have correct dtype."""
    # Create instance and set dtype
    aug = RandomThinPlateSpline(scale=0.2, p=1.0)
    aug.set_rng_device_and_dtype(torch.device("cpu"), torch.float32)
    
    # Generate parameters
    shape = (2, 3, 32, 32)
    codeflash_output = aug.generate_parameters(shape); params = codeflash_output # 72.1μs -> 64.2μs (12.3% faster)

def test_generate_parameters_device_placement():
    """Test that output tensors are on correct device."""
    # Create instance
    aug = RandomThinPlateSpline(scale=0.2, p=1.0)
    aug.set_rng_device_and_dtype(torch.device("cpu"), torch.float32)
    
    # Generate parameters
    shape = (2, 3, 32, 32)
    codeflash_output = aug.generate_parameters(shape); params = codeflash_output # 72.2μs -> 64.3μs (12.3% faster)

def test_generate_parameters_different_padding_modes():
    """Test that different padding modes don't affect parameter generation."""
    # Test with different padding modes
    padding_modes = [
        SamplePadding.ZEROS,
        SamplePadding.BORDER,
        SamplePadding.REFLECTION,
        "zeros",
        0
    ]
    
    shape = (2, 3, 32, 32)
    
    for padding_mode in padding_modes:
        aug = RandomThinPlateSpline(scale=0.2, padding_mode=padding_mode, p=1.0)
        codeflash_output = aug.generate_parameters(shape); params = codeflash_output # 261μs -> 221μs (18.1% faster)

def test_generate_parameters_multiple_calls_different_results():
    """Test that multiple calls produce different results (randomness)."""
    # Create instance with same_on_batch=False
    aug = RandomThinPlateSpline(scale=0.2, same_on_batch=False, p=1.0)
    
    # Generate parameters twice
    shape = (3, 3, 32, 32)
    codeflash_output = aug.generate_parameters(shape); params1 = codeflash_output # 71.0μs -> 64.7μs (9.67% faster)
    codeflash_output = aug.generate_parameters(shape); params2 = codeflash_output # 42.1μs -> 32.8μs (28.6% faster)

# Test Large Scale Cases
def test_generate_parameters_large_batch_size():
    """Test with large batch size to assess scalability."""
    # Create instance
    aug = RandomThinPlateSpline(scale=0.2, p=1.0)
    
    # Test with large batch size (but not too large to avoid memory issues)
    shape = (128, 3, 32, 32)
    codeflash_output = aug.generate_parameters(shape); params = codeflash_output # 102μs -> 75.5μs (35.9% faster)
    
    # Verify noise bounds
    noise = params["dst"] - params["src"]

def test_generate_parameters_large_batch_same_on_batch():
    """Test large batch with same_on_batch=True."""
    # Create instance with same_on_batch=True
    aug = RandomThinPlateSpline(scale=0.2, same_on_batch=True, p=1.0)
    
    # Test with large batch
    shape = (100, 3, 64, 64)
    codeflash_output = aug.generate_parameters(shape); params = codeflash_output # 81.5μs -> 73.9μs (10.2% faster)
    
    # Calculate noise
    noise = params["dst"] - params["src"]
    
    # Verify all batch items have identical noise
    for i in range(1, 100):
        pass

def test_generate_parameters_stress_test_repeated_calls():
    """Stress test with many repeated calls."""
    # Create instance
    aug = RandomThinPlateSpline(scale=0.2, p=1.0)
    
    # Make many calls
    shape = (4, 3, 32, 32)
    for _ in range(100):
        codeflash_output = aug.generate_parameters(shape); params = codeflash_output # 3.55ms -> 2.76ms (28.8% faster)
        
        noise = params["dst"] - params["src"]

def test_generate_parameters_various_scales():
    """Test with various scale values to ensure robustness."""
    # Test with different scale values
    scales = [0.0, 0.1, 0.5, 1.0, 2.0, 10.0]
    shape = (3, 3, 32, 32)
    
    for scale in scales:
        aug = RandomThinPlateSpline(scale=scale, p=1.0)
        codeflash_output = aug.generate_parameters(shape); params = codeflash_output # 303μs -> 269μs (12.8% faster)
        
        # Calculate noise
        noise = params["dst"] - params["src"]

def test_generate_parameters_consistency_with_fixed_seed():
    """Test that results are consistent when using the same random seed."""
    # Set random seed
    torch.manual_seed(42)
    
    # Create instance and generate parameters
    aug1 = RandomThinPlateSpline(scale=0.2, p=1.0)
    shape = (3, 3, 32, 32)
    codeflash_output = aug1.generate_parameters(shape); params1 = codeflash_output # 72.6μs -> 77.7μs (6.60% slower)
    
    # Reset seed and repeat
    torch.manual_seed(42)
    aug2 = RandomThinPlateSpline(scale=0.2, p=1.0)
    codeflash_output = aug2.generate_parameters(shape); params2 = codeflash_output # 67.8μs -> 57.8μs (17.3% faster)

def test_generate_parameters_tensor_contiguity():
    """Test that generated tensors are contiguous in memory."""
    # Create instance
    aug = RandomThinPlateSpline(scale=0.2, p=1.0)
    
    # Generate parameters
    shape = (4, 3, 32, 32)
    codeflash_output = aug.generate_parameters(shape); params = codeflash_output # 74.2μs -> 64.7μs (14.8% faster)

def test_generate_parameters_no_nan_or_inf():
    """Test that generated parameters don't contain NaN or Inf values."""
    # Create instance
    aug = RandomThinPlateSpline(scale=0.2, p=1.0)
    
    # Generate parameters
    shape = (5, 3, 32, 32)
    codeflash_output = aug.generate_parameters(shape); params = codeflash_output # 73.3μs -> 74.5μs (1.58% slower)

def test_generate_parameters_align_corners_flag():
    """Test that align_corners flag doesn't affect parameter generation."""
    # Test with different align_corners values
    for align_corners in [True, False]:
        aug = RandomThinPlateSpline(scale=0.2, align_corners=align_corners, p=1.0)
        shape = (2, 3, 32, 32)
        codeflash_output = aug.generate_parameters(shape); params = codeflash_output # 123μs -> 106μs (16.2% 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-RandomThinPlateSpline.generate_parameters-mkdtfmha and push.

[github-actions]

⚠️ PR Validation Warnings

❌ No linked issue found: This PR does not reference any issue. Please link to an issue using "Fixes kornia#123" or "Closes kornia#123" in the PR description.

---

Note: This PR can remain open, but please address these issues to ensure a smooth review process. For more information, see our Contributing Guide.

[github-actions]

This pull request has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs within 7 days. Thank you for your contributions!

[github-actions]

This pull request has been automatically closed due to inactivity. Feel free to reopen it if you would like to continue working on it.