⚡️ Speed up function batched_squared_norm by 19%

[codeflash-ai]

📄 19% (0.19x) speedup for batched_squared_norm in kornia/geometry/linalg.py

⏱️ Runtime : 279 microseconds → 235 microseconds (best of 219 runs)

📝 Explanation and details

Key Changes.

1. Optimization in batched_dot_product.

   • Kept the function structure the same to maintain compatibility.
   • Use x.mul(y) instead of x * y to make it explicit and potentially more optimized within Tensor libraries that have optimized mul implementations.

2. Optimization in batched_squared_norm.

   • Used x.square() to directly compute the element-wise squares more efficiently, and immediately followed by summation with sum(dim=-1, keepdim=keepdim). This reduces the number of dot product operations and leverages faster library-level operations.

By focusing on using efficient tensor operations, these changes ensure faster execution times by reducing overhead and leveraging underlying optimized mathematical operations provided by the tensor library being used.

✅ Correctness verification report:

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

🌀 Generated Regression Tests Details

from __future__ import annotations

from typing import TypeVar

# imports
import pytest  # used for our unit tests
import torch  # used to create tensors for testing
from kornia.core import Tensor
from kornia.core.check import KORNIA_CHECK_SHAPE
from kornia.geometry.linalg import batched_squared_norm

# function to test
# LICENSE HEADER MANAGED BY add-license-header
#
# Copyright 2018 Kornia Team
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#


__all__ = [
    "batched_dot_product",
    "batched_squared_norm",
    "compose_transformations",
    "euclidean_distance",
    "inverse_transformation",
    "point_line_distance",
    "relative_transformation",
    "squared_norm",
    "transform_points",
]
from kornia.geometry.linalg import batched_squared_norm

# aliases
squared_norm = batched_squared_norm

# TODO:
# - project_points: from opencv

# LICENSE HEADER MANAGED BY add-license-header
#
# Copyright 2018 Kornia Team
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""The testing package contains testing-specific utilities."""

__all__ = [
    "KORNIA_CHECK",
    "KORNIA_CHECK_DM_DESC",
    "KORNIA_CHECK_IS_COLOR",
    "KORNIA_CHECK_IS_GRAY",
    "KORNIA_CHECK_IS_IMAGE",
    "KORNIA_CHECK_IS_LIST_OF_TENSOR",
    "KORNIA_CHECK_IS_TENSOR",
    "KORNIA_CHECK_LAF",
    "KORNIA_CHECK_SAME_DEVICE",
    "KORNIA_CHECK_SAME_DEVICES",
    "KORNIA_CHECK_SHAPE",
    "KORNIA_CHECK_TYPE",
    "KORNIA_UNWRAP",
]

# Logger api


# TODO: add somehow type check, or enforce to do it before
def KORNIA_CHECK_SHAPE(x: Tensor, shape: list[str], raises: bool = True) -> bool:
    """Check whether a tensor has a specified shape.

    The shape can be specified with a implicit or explicit list of strings.
    The guard also check whether the variable is a type `Tensor`.

    Args:
        x: the tensor to evaluate.
        shape: a list with strings with the expected shape.
        raises: bool indicating whether an exception should be raised upon failure.

    Raises:
        Exception: if the input tensor is has not the expected shape and raises is True.

    Example:
        >>> x = torch.rand(2, 3, 4, 4)
        >>> KORNIA_CHECK_SHAPE(x, ["B", "C", "H", "W"])  # implicit
        True

        >>> x = torch.rand(2, 3, 4, 4)
        >>> KORNIA_CHECK_SHAPE(x, ["2", "3", "H", "W"])  # explicit
        True

    """
    if "*" == shape[0]:
        shape_to_check = shape[1:]
        x_shape_to_check = x.shape[-len(shape) + 1 :]
    elif "*" == shape[-1]:
        shape_to_check = shape[:-1]
        x_shape_to_check = x.shape[: len(shape) - 1]
    else:
        shape_to_check = shape
        x_shape_to_check = x.shape

    if len(x_shape_to_check) != len(shape_to_check):
        if raises:
            raise TypeError(f"{x} shape must be [{shape}]. Got {x.shape}")
        else:
            return False

    for i in range(len(x_shape_to_check)):
        # The voodoo below is because torchscript does not like
        # that dim can be both int and str
        dim_: str = shape_to_check[i]
        if not dim_.isnumeric():
            continue
        dim = int(dim_)
        if x_shape_to_check[i] != dim:
            if raises:
                raise TypeError(f"{x} shape must be [{shape}]. Got {x.shape}")
            else:
                return False
    return True


T = TypeVar("T", bound=type)


# unit tests

def test_single_vector():
    # Test with a simple 1D tensor
    x = torch.tensor([1.0, 2.0, 3.0])
    expected = 14.0  # 1^2 + 2^2 + 3^2 = 14

def test_batch_of_vectors():
    # Test with a 2D tensor representing a batch of vectors
    x = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
    expected = torch.tensor([5.0, 25.0])  # [1^2 + 2^2, 3^2 + 4^2]


def test_single_element_tensor():
    # Test with a tensor with a single element
    x = torch.tensor([5.0])
    expected = 25.0  # 5^2 = 25

def test_high_dimensionality():
    # Test with a higher dimensional tensor
    x = torch.tensor([[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]])
    expected = torch.tensor([[5.0, 25.0], [61.0, 113.0]])

def test_zero_vector():
    # Test with a tensor of zeros
    x = torch.tensor([0.0, 0.0, 0.0])
    expected = 0.0

def test_inf_values():
    # Test with a tensor containing inf
    x = torch.tensor([float('inf'), 1.0, 2.0])
    expected = float('inf')

def test_nan_values():
    # Test with a tensor containing NaN
    x = torch.tensor([float('nan'), 1.0, 2.0])
    codeflash_output = batched_squared_norm(x)

def test_keepdim_true():
    # Test with keepdim=True
    x = torch.tensor([1.0, 2.0, 3.0])
    expected = torch.tensor([14.0])

def test_large_tensor():
    # Test with a large tensor
    x = torch.rand(1000)
    codeflash_output = batched_squared_norm(x)

def test_large_batch_of_vectors():
    # Test with a large batch of vectors
    x = torch.rand(100, 100)
    codeflash_output = batched_squared_norm(x)



def test_mixed_data_types():
    # Test with mixed data types
    x = torch.tensor([1, 2.5, 3])
    expected = 1**2 + 2.5**2 + 3**2

def test_multi_device_support():
    # Test with tensors on different devices
    if torch.cuda.is_available():
        x = torch.tensor([1.0, 2.0, 3.0]).cuda()
        expected = 14.0
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from __future__ import annotations

from typing import TypeVar

# imports
import pytest  # used for our unit tests
import torch  # used for creating tensors
from kornia.core import Tensor
from kornia.core.check import KORNIA_CHECK_SHAPE
from kornia.geometry.linalg import batched_squared_norm

# function to test
# LICENSE HEADER MANAGED BY add-license-header
#
# Copyright 2018 Kornia Team
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#



__all__ = [
    "batched_dot_product",
    "batched_squared_norm",
    "compose_transformations",
    "euclidean_distance",
    "inverse_transformation",
    "point_line_distance",
    "relative_transformation",
    "squared_norm",
    "transform_points",
]
from kornia.geometry.linalg import batched_squared_norm

# aliases
squared_norm = batched_squared_norm

# TODO:
# - project_points: from opencv

# LICENSE HEADER MANAGED BY add-license-header
#
# Copyright 2018 Kornia Team
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""The testing package contains testing-specific utilities."""

__all__ = [
    "KORNIA_CHECK",
    "KORNIA_CHECK_DM_DESC",
    "KORNIA_CHECK_IS_COLOR",
    "KORNIA_CHECK_IS_GRAY",
    "KORNIA_CHECK_IS_IMAGE",
    "KORNIA_CHECK_IS_LIST_OF_TENSOR",
    "KORNIA_CHECK_IS_TENSOR",
    "KORNIA_CHECK_LAF",
    "KORNIA_CHECK_SAME_DEVICE",
    "KORNIA_CHECK_SAME_DEVICES",
    "KORNIA_CHECK_SHAPE",
    "KORNIA_CHECK_TYPE",
    "KORNIA_UNWRAP",
]

# Logger api


# TODO: add somehow type check, or enforce to do it before
def KORNIA_CHECK_SHAPE(x: Tensor, shape: list[str], raises: bool = True) -> bool:
    """Check whether a tensor has a specified shape.

    The shape can be specified with a implicit or explicit list of strings.
    The guard also check whether the variable is a type `Tensor`.

    Args:
        x: the tensor to evaluate.
        shape: a list with strings with the expected shape.
        raises: bool indicating whether an exception should be raised upon failure.

    Raises:
        Exception: if the input tensor is has not the expected shape and raises is True.

    Example:
        >>> x = torch.rand(2, 3, 4, 4)
        >>> KORNIA_CHECK_SHAPE(x, ["B", "C", "H", "W"])  # implicit
        True

        >>> x = torch.rand(2, 3, 4, 4)
        >>> KORNIA_CHECK_SHAPE(x, ["2", "3", "H", "W"])  # explicit
        True

    """
    if "*" == shape[0]:
        shape_to_check = shape[1:]
        x_shape_to_check = x.shape[-len(shape) + 1 :]
    elif "*" == shape[-1]:
        shape_to_check = shape[:-1]
        x_shape_to_check = x.shape[: len(shape) - 1]
    else:
        shape_to_check = shape
        x_shape_to_check = x.shape

    if len(x_shape_to_check) != len(shape_to_check):
        if raises:
            raise TypeError(f"{x} shape must be [{shape}]. Got {x.shape}")
        else:
            return False

    for i in range(len(x_shape_to_check)):
        # The voodoo below is because torchscript does not like
        # that dim can be both int and str
        dim_: str = shape_to_check[i]
        if not dim_.isnumeric():
            continue
        dim = int(dim_)
        if x_shape_to_check[i] != dim:
            if raises:
                raise TypeError(f"{x} shape must be [{shape}]. Got {x.shape}")
            else:
                return False
    return True


T = TypeVar("T", bound=type)


# unit tests

def test_single_vector():
    # Test with a single 1D tensor
    x = torch.tensor([3.0, 4.0])
    expected = 25.0
    codeflash_output = batched_squared_norm(x)

def test_single_vector_negative():
    # Test with a single 1D tensor with negative values
    x = torch.tensor([-3.0, -4.0])
    expected = 25.0
    codeflash_output = batched_squared_norm(x)

def test_batch_of_vectors():
    # Test with a 2D tensor with multiple vectors
    x = torch.tensor([[3.0, 4.0], [1.0, 2.0]])
    expected = torch.tensor([25.0, 5.0])
    codeflash_output = batched_squared_norm(x)

def test_zero_vector():
    # Test with a tensor of zeros
    x = torch.tensor([0.0, 0.0])
    expected = 0.0
    codeflash_output = batched_squared_norm(x)

def test_batch_of_zero_vectors():
    # Test with a batch of zero vectors
    x = torch.tensor([[0.0, 0.0], [0.0, 0.0]])
    expected = torch.tensor([0.0, 0.0])
    codeflash_output = batched_squared_norm(x)


def test_high_dimensionality():
    # Test with a tensor with more than two dimensions
    x = torch.tensor([[[3.0, 4.0], [1.0, 2.0]]])
    expected = torch.tensor([[25.0, 5.0]])
    codeflash_output = batched_squared_norm(x, keepdim=True)

def test_special_values():
    # Test with a tensor containing inf or nan
    x = torch.tensor([float('inf'), 1.0])
    codeflash_output = batched_squared_norm(x)

def test_large_scale():
    # Test with a very large tensor
    x = torch.rand(1000, 100)  # Ensure size is under 100MB
    codeflash_output = batched_squared_norm(x)

def test_keepdim():
    # Test with keepdim=True
    x = torch.tensor([[3.0, 4.0]])
    expected = torch.tensor([[25.0]])
    codeflash_output = batched_squared_norm(x, keepdim=True)

To edit these changes git checkout codeflash/optimize-batched_squared_norm-m8oca7qu 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!