Enable torch compile on _compute_affine_output_size

Fixed failing tests:
test.test_transforms_v2.TestRotate.test_kernel_image[cpu-dtype0-expand-True]
test.test_transforms_v2.TestRotate.test_kernel_image[cpu-dtype1-expand-True]
test.test_transforms_v2.TestRotate.test_kernel_bounding_boxes[cpu-dtype0-BoundingBoxFormat.XYXY-expand-True]
test.test_transforms_v2.TestRotate.test_kernel_bounding_boxes[cpu-dtype0-BoundingBoxFormat.XYWH-expand-True]
test.test_transforms_v2.TestRotate.test_kernel_bounding_boxes[cpu-dtype0-BoundingBoxFormat.CXCYWH-expand-True]
test.test_transforms_v2.TestRotate.test_kernel_bounding_boxes[cpu-dtype1-BoundingBoxFormat.XYXY-expand-True]
test.test_transforms_v2.TestRotate.test_kernel_bounding_boxes[cpu-dtype1-BoundingBoxFormat.XYWH-expand-True]
test.test_transforms_v2.TestRotate.test_kernel_bounding_boxes[cpu-dtype1-BoundingBoxFormat.CXCYWH-expand-True]

from compile (true, inductor, true) / dynamic=true,backend=inductor,fullgraph=true summary
https://github.com/pytorch/vision/actions/runs/7298953575

torchvision/transforms/v2/functional/_geometry.py

@@ -525,6 +525,13 @@ def _get_inverse_affine_matrix(
 
 
 def _compute_affine_output_size(matrix: List[float], w: int, h: int) -> Tuple[int, int]:
+    if torch._dynamo.is_compiling():
+        return _compute_affine_output_size_python(matrix, w, h)
+    else:
+        return _compute_affine_output_size_tensor(matrix, w, h)
+
+
+def _compute_affine_output_size_tensor(matrix: List[float], w: int, h: int) -> Tuple[int, int]:
     # Inspired of PIL implementation:
     # https://github.com/python-pillow/Pillow/blob/11de3318867e4398057373ee9f12dcb33db7335c/src/PIL/Image.py#L2054
 
@@ -559,6 +566,58 @@ def _compute_affine_output_size(matrix: List[float], w: int, h: int) -> Tuple[in
     return int(size[0]), int(size[1])  # w, h
 
 
+def _compute_affine_output_size_python(matrix: List[float], w: int, h: int) -> Tuple[int, int]:
+    # Same method as _compute_affine_output_size_tensor but not using torch tensor and used for torch compile
+
+    # Inspired of PIL implementation:
+    # https://github.com/python-pillow/Pillow/blob/11de3318867e4398057373ee9f12dcb33db7335c/src/PIL/Image.py#L2054
+
+    # pts are Top-Left, Top-Right, Bottom-Left, Bottom-Right points.
+    # Points are shifted due to affine matrix torch convention about
+    # the center point. Center is (0, 0) for image center pivot point (w * 0.5, h * 0.5)
+    half_w = 0.5 * w
+    half_h = 0.5 * h
+
+    m0, m1, m2, m3, m4, m5 = matrix
+
+    m0_half_w = m0 * half_w
+    m1_half_h = m1 * half_h
+    m3_half_w = m3 * half_w
+    m4_half_h = m4 * half_h
+
+    new_pts_x = [
+        m2 - m0_half_w - m1_half_h,
+        m2 - m0_half_w + m1_half_h,
+        m2 + m0_half_w + m1_half_h,
+        m2 + m0_half_w - m1_half_h,
+    ]
+
+    new_pts_y = [
+        m5 - m3_half_w - m4_half_h,
+        m5 - m3_half_w + m4_half_h,
+        m5 + m3_half_w + m4_half_h,
+        m5 + m3_half_w - m4_half_h,
+    ]
+
+    min_val_x = min(new_pts_x) + half_w
+    max_val_x = max(new_pts_x) + half_w
+    min_val_y = min(new_pts_y) + half_h
+    max_val_y = max(new_pts_y) + half_h
+
+    # Truncate precision to 1e-4 to avoid ceil of Xe-15 to 1.0
+    tol = 1e-4
+    inv_tol = 1.0 / tol
+
+    cmax_x = math.ceil(int(max_val_x * inv_tol) * tol)
+    cmax_y = math.ceil(int(max_val_y * inv_tol) * tol)
+    cmin_x = math.floor(int(min_val_x * inv_tol) * tol)
+    cmin_y = math.floor(int(min_val_y * inv_tol) * tol)
+
+    size_x = cmax_x - cmin_x
+    size_y = cmax_y - cmin_y
+    return int(size_x), int(size_y)  # w, h
+
+
 def _apply_grid_transform(img: torch.Tensor, grid: torch.Tensor, mode: str, fill: _FillTypeJIT) -> torch.Tensor:
     input_shape = img.shape
     output_height, output_width = grid.shape[1], grid.shape[2]