⚡️ Speed up function _compute_shear_matrix by 14%

kornia/geometry/transform/affwarp.py

@@ -122,11 +122,16 @@ def _compute_scaling_matrix(scale: Tensor, center: Tensor) -> Tensor:
 
 def _compute_shear_matrix(shear: Tensor) -> Tensor:
     """Compute affine matrix for shearing."""
+    # Create a 3x3 identity like matrix using the optimized eye_like function
     matrix: Tensor = eye_like(3, shear, shared_memory=False)
 
+    # Split the shear tensor into two separate tensors for x and y shearing
     shx, shy = torch.chunk(shear, chunks=2, dim=-1)
-    matrix[..., 0, 1:2] += shx
-    matrix[..., 1, 0:1] += shy
+
+    # Efficiently adding shearing values to the identity matrix to form the resulting shear matrix
+    matrix[..., 0, 1] = shx.squeeze(-1)
+    matrix[..., 1, 0] = shy.squeeze(-1)
+
     return matrix
 
 

kornia/utils/misc.py

@@ -45,9 +45,14 @@ def eye_like(n: int, input: Tensor, shared_memory: bool = False) -> Tensor:
     if len(input.shape) < 1:
         raise AssertionError(input.shape)
 
+    # Create identity matrix of size n
     identity = eye(n, device=input.device).type(input.dtype)
 
-    return identity[None].expand(input.shape[0], n, n) if shared_memory else identity[None].repeat(input.shape[0], 1, 1)
+    # Use expand if shared_memory is True; otherwise use repeat
+    if shared_memory:
+        return identity[None].expand(input.shape[0], n, n)
+    else:
+        return identity.repeat([input.shape[0], 1, 1])
 
 
 def vec_like(n: int, tensor: Tensor, shared_memory: bool = False) -> Tensor:

pyproject.toml

@@ -268,3 +268,11 @@ ignore_errors = true
 
 [tool.pydocstyle]
 match = '.*\.py'
+
+[tool.codeflash]
+# All paths are relative to this pyproject.toml's directory.
+module-root = "kornia"
+tests-root = "tests"
+test-framework = "pytest"
+ignore-paths = []
+formatter-cmds = ["ruff check --exit-zero --fix $file", "ruff format $file"]