⚡️ Speed up function rotate3d by 6%

[codeflash-ai]

📄 6% (0.06x) speedup for rotate3d in kornia/geometry/transform/affwarp.py

⏱️ Runtime : 63.4 milliseconds → 59.8 milliseconds (best of 100 runs)

📝 Explanation and details

Optimization Summary.

1. Eliminated redundant operations by consolidating tensor manipulations and matrix forming processes to array operations.
2. Reduced usage of unsqueeze and expand by strategically restructuring tensor operations and conditions.
3. Minimized intermediate calculations and unnecessary assignments in _compute_rotation_matrix3d and affine3d operations. This allowed for more direct tensor construction and transformation.
4. The logic in rotate3d and affine3d functions has been optimized to minimize tensor shape checks and unnecessary reshaping, allowing for more efficient tensor manipulations.

✅ Correctness verification report:

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

🌀 Generated Regression Tests Details

import pytest
import torch
from kornia.core import tensor
from kornia.geometry.transform.affwarp import rotate3d

# Basic Functionality Tests

def test_single_axis_rotation_yaw():
    # Create a simple tensor with shape (1, 1, 3, 3, 3)
    input_tensor = torch.arange(27, dtype=torch.float32).reshape(1, 1, 3, 3, 3)
    # Rotate 90 degrees around yaw axis
    yaw = torch.tensor([90.0])
    pitch = torch.tensor([0.0])
    roll = torch.tensor([0.0])
    codeflash_output = rotate3d(input_tensor, yaw, pitch, roll)

def test_combined_axis_rotation():
    input_tensor = torch.arange(27, dtype=torch.float32).reshape(1, 1, 3, 3, 3)
    yaw = torch.tensor([45.0])
    pitch = torch.tensor([45.0])
    roll = torch.tensor([45.0])
    codeflash_output = rotate3d(input_tensor, yaw, pitch, roll)

# Edge Cases

def test_zero_rotation():
    input_tensor = torch.arange(27, dtype=torch.float32).reshape(1, 1, 3, 3, 3)
    yaw = torch.tensor([0.0])
    pitch = torch.tensor([0.0])
    roll = torch.tensor([0.0])
    codeflash_output = rotate3d(input_tensor, yaw, pitch, roll)

def test_negative_angles():
    input_tensor = torch.arange(27, dtype=torch.float32).reshape(1, 1, 3, 3, 3)
    yaw = torch.tensor([-90.0])
    pitch = torch.tensor([-45.0])
    roll = torch.tensor([-30.0])
    codeflash_output = rotate3d(input_tensor, yaw, pitch, roll)

# Input Validation

def test_invalid_tensor_type():
    with pytest.raises(TypeError):
        rotate3d([1, 2, 3], torch.tensor([0.0]), torch.tensor([0.0]), torch.tensor([0.0]))

def test_invalid_angle_type():
    with pytest.raises(TypeError):
        rotate3d(torch.rand(1, 1, 3, 3, 3), [0.0], torch.tensor([0.0]), torch.tensor([0.0]))

def test_invalid_tensor_shape():
    with pytest.raises(ValueError):
        rotate3d(torch.rand(1, 1, 1, 1, 1, 1), torch.tensor([0.0]), torch.tensor([0.0]), torch.tensor([0.0]))

# Performance and Scalability

def test_large_scale_rotation():
    # Create a large tensor with shape (1, 1, 64, 64, 64) which is under 100MB
    input_tensor = torch.rand(1, 1, 64, 64, 64)
    yaw = torch.tensor([45.0])
    pitch = torch.tensor([45.0])
    roll = torch.tensor([45.0])
    codeflash_output = rotate3d(input_tensor, yaw, pitch, roll)

# Special Cases

def test_non_centered_rotation():
    input_tensor = torch.arange(27, dtype=torch.float32).reshape(1, 1, 3, 3, 3)
    yaw = torch.tensor([45.0])
    pitch = torch.tensor([45.0])
    roll = torch.tensor([45.0])
    center = torch.tensor([[1.0, 1.0, 1.0]])
    codeflash_output = rotate3d(input_tensor, yaw, pitch, roll, center=center)

def test_batch_processing():
    input_tensor = torch.arange(54, dtype=torch.float32).reshape(2, 1, 3, 3, 3)
    yaw = torch.tensor([45.0, 90.0])
    pitch = torch.tensor([45.0, 0.0])
    roll = torch.tensor([45.0, 0.0])
    codeflash_output = rotate3d(input_tensor, yaw, pitch, roll)

# Interpolation and Padding Modes

def test_different_interpolation_modes():
    input_tensor = torch.arange(27, dtype=torch.float32).reshape(1, 1, 3, 3, 3)
    yaw = torch.tensor([45.0])
    pitch = torch.tensor([45.0])
    roll = torch.tensor([45.0])
    codeflash_output = rotate3d(input_tensor, yaw, pitch, roll, mode='bilinear')
    codeflash_output = rotate3d(input_tensor, yaw, pitch, roll, mode='nearest')

def test_different_padding_modes():
    input_tensor = torch.arange(27, dtype=torch.float32).reshape(1, 1, 3, 3, 3)
    yaw = torch.tensor([45.0])
    pitch = torch.tensor([45.0])
    roll = torch.tensor([45.0])
    codeflash_output = rotate3d(input_tensor, yaw, pitch, roll, padding_mode='zeros')
    codeflash_output = rotate3d(input_tensor, yaw, pitch, roll, padding_mode='border')
    codeflash_output = rotate3d(input_tensor, yaw, pitch, roll, padding_mode='reflection')

# Alignment Corner Cases

def test_align_corners_flag():
    input_tensor = torch.arange(27, dtype=torch.float32).reshape(1, 1, 3, 3, 3)
    yaw = torch.tensor([45.0])
    pitch = torch.tensor([45.0])
    roll = torch.tensor([45.0])
    codeflash_output = rotate3d(input_tensor, yaw, pitch, roll, align_corners=True)
    codeflash_output = rotate3d(input_tensor, yaw, pitch, roll, align_corners=False)
# 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.geometry.conversions import (axis_angle_to_rotation_matrix,
                                         convert_affinematrix_to_homography3d,
                                         deg2rad, normalize_homography3d)
from kornia.geometry.transform.affwarp import rotate3d
from kornia.geometry.transform.imgwarp import warp_affine3d
from kornia.utils import eye_like
from kornia.utils.helpers import _torch_inverse_cast
from torch import Tensor

# function to test (rotate3d and its dependencies are already defined above)

# unit tests

def test_rotate3d_single_rotation_yaw():
    """Test rotating only with yaw."""
    tensor = torch.rand(1, 1, 10, 10, 10)  # Create a random tensor
    yaw = torch.tensor([90.0])  # 90 degrees yaw
    pitch = torch.tensor([0.0])
    roll = torch.tensor([0.0])
    codeflash_output = rotate3d(tensor, yaw, pitch, roll)

def test_rotate3d_single_rotation_pitch():
    """Test rotating only with pitch."""
    tensor = torch.rand(1, 1, 10, 10, 10)  # Create a random tensor
    yaw = torch.tensor([0.0])
    pitch = torch.tensor([45.0])  # 45 degrees pitch
    roll = torch.tensor([0.0])
    codeflash_output = rotate3d(tensor, yaw, pitch, roll)

def test_rotate3d_single_rotation_roll():
    """Test rotating only with roll."""
    tensor = torch.rand(1, 1, 10, 10, 10)  # Create a random tensor
    yaw = torch.tensor([0.0])
    pitch = torch.tensor([0.0])
    roll = torch.tensor([30.0])  # 30 degrees roll
    codeflash_output = rotate3d(tensor, yaw, pitch, roll)

def test_rotate3d_multiple_rotations():
    """Test rotating with yaw, pitch, and roll."""
    tensor = torch.rand(1, 1, 10, 10, 10)  # Create a random tensor
    yaw = torch.tensor([45.0])
    pitch = torch.tensor([45.0])
    roll = torch.tensor([45.0])
    codeflash_output = rotate3d(tensor, yaw, pitch, roll)

def test_rotate3d_zero_rotation():
    """Test rotating with zero angles."""
    tensor = torch.rand(1, 1, 10, 10, 10)  # Create a random tensor
    yaw = torch.tensor([0.0])
    pitch = torch.tensor([0.0])
    roll = torch.tensor([0.0])
    codeflash_output = rotate3d(tensor, yaw, pitch, roll)

def test_rotate3d_negative_angles():
    """Test rotating with negative angles."""
    tensor = torch.rand(1, 1, 10, 10, 10)  # Create a random tensor
    yaw = torch.tensor([-45.0])
    pitch = torch.tensor([0.0])
    roll = torch.tensor([0.0])
    codeflash_output = rotate3d(tensor, yaw, pitch, roll)

def test_rotate3d_large_angles():
    """Test rotating with angles greater than 360 degrees."""
    tensor = torch.rand(1, 1, 10, 10, 10)  # Create a random tensor
    yaw = torch.tensor([720.0])  # 720 degrees yaw
    pitch = torch.tensor([0.0])
    roll = torch.tensor([0.0])
    codeflash_output = rotate3d(tensor, yaw, pitch, roll)

def test_rotate3d_different_tensor_shapes():
    """Test rotating tensors of different shapes."""
    tensor = torch.rand(1, 3, 20, 20, 20)  # Create a random tensor
    yaw = torch.tensor([90.0])
    pitch = torch.tensor([0.0])
    roll = torch.tensor([0.0])
    codeflash_output = rotate3d(tensor, yaw, pitch, roll)

def test_rotate3d_custom_center():
    """Test rotating with a specified center."""
    tensor = torch.rand(1, 1, 10, 10, 10)  # Create a random tensor
    yaw = torch.tensor([90.0])
    pitch = torch.tensor([0.0])
    roll = torch.tensor([0.0])
    center = torch.tensor([[5.0, 5.0, 5.0]])  # Center at (5, 5, 5)
    codeflash_output = rotate3d(tensor, yaw, pitch, roll, center)

def test_rotate3d_interpolation_modes():
    """Test rotating with different interpolation modes."""
    tensor = torch.rand(1, 1, 10, 10, 10)  # Create a random tensor
    yaw = torch.tensor([90.0])
    pitch = torch.tensor([0.0])
    roll = torch.tensor([0.0])
    codeflash_output = rotate3d(tensor, yaw, pitch, roll, mode='bilinear')
    codeflash_output = rotate3d(tensor, yaw, pitch, roll, mode='nearest')

def test_rotate3d_large_tensor():
    """Test rotating a large tensor to assess performance."""
    tensor = torch.rand(1, 1, 100, 100, 100)  # Create a large tensor
    yaw = torch.tensor([45.0])
    pitch = torch.tensor([45.0])
    roll = torch.tensor([45.0])
    codeflash_output = rotate3d(tensor, yaw, pitch, roll)

def test_rotate3d_invalid_types():
    """Test handling of invalid input types."""
    tensor = torch.rand(1, 1, 10, 10, 10)  # Create a random tensor
    yaw = [90.0]  # Invalid type: list instead of tensor
    pitch = torch.tensor([0.0])
    roll = torch.tensor([0.0])
    with pytest.raises(TypeError):
        rotate3d(tensor, yaw, pitch, roll)

def test_rotate3d_invalid_shapes():
    """Test handling of invalid tensor shapes."""
    tensor = torch.rand(10, 10, 10)  # Invalid shape: missing batch and channel dimensions
    yaw = torch.tensor([90.0])
    pitch = torch.tensor([0.0])
    roll = torch.tensor([0.0])
    with pytest.raises(ValueError):
        rotate3d(tensor, yaw, pitch, roll)
# 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-rotate3d-m8o9mb6t and push.

[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!