Resolve Co-pilot review comments

Author: mansnils

backends/xnnpack/operators/op_mean_dim.py

@@ -58,6 +58,13 @@ def define_node(
             len(input_shape) == 4, "Require input to mean.dim be 4 dimensional"
         )
 
+        # This visitor serializes mean.dim as Global Average Pooling, which has
+        # no field for an explicit dtype override.
+        check_or_raise(
+            node.kwargs.get("dtype") is None,
+            "XNNPACK does not support mean.dim with dtype",
+        )
+
         # mean dims
         mean_dims = normalize_mean_dims(node.args[1], len(input_shape))
         check_or_raise(

backends/xnnpack/partition/config/generic_node_configs.py

@@ -517,20 +517,30 @@ def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
             return False
 
         input_rank = get_input_node(node, 0).meta["val"].dim()
-        keepdim = len(node.args) >= 3 and bool(node.args[2])
-        dims = normalize_mean_dims(node.args[1], input_rank)
-
-        if sorted(dims) != [2, 3]:
+        if input_rank != 4:
             why(
                 node,
-                reason="mean.dim only supports averaging 4D tensors across the innermost dimensions",
+                reason=f"mean.dim only supports averaging 4D tensors, got tensor of rank {input_rank}",
             )
             return False
 
-        if input_rank != 4:
+        # This path lowers mean.dim to XNNPACK Global Average Pooling, which
+        # cannot encode an explicit dtype override.
+        if node.kwargs.get("dtype") is not None:
+            why(node, reason="mean.dim does not support dtype")
+            return False
+
+        keepdim = len(node.args) >= 3 and bool(node.args[2])
+        try:
+            dims = normalize_mean_dims(node.args[1], input_rank)
+        except ValueError as error:
+            why(node, reason=f"mean.dim has invalid dims: {error}")
+            return False
+
+        if sorted(dims) != [2, 3]:
             why(
                 node,
-                reason=f"mean.dim only supports averaging 4D tensors, got tensor of rank {input_rank}",
+                reason="mean.dim only supports averaging 4D tensors across the innermost dimensions",
             )
             return False
 

backends/xnnpack/test/ops/test_mean_dim.py

@@ -16,14 +16,18 @@ def setUp(self):
         torch._dynamo.reset()
 
     class MeanDim(torch.nn.Module):
-        def __init__(self, dims, keepdim=True):
+        def __init__(self, dims, keepdim=True, dtype=None):
             super().__init__()
             self.dims = dims
             self.keepdim = keepdim
+            self.dtype = dtype
 
         def forward(self, x):
             y = x + x
-            z = torch.mean(y, self.dims, keepdim=self.keepdim)
+            if self.dtype is None:
+                z = torch.mean(y, self.dims, keepdim=self.keepdim)
+            else:
+                z = torch.mean(y, self.dims, keepdim=self.keepdim, dtype=self.dtype)
             return z
 
     def _test_mean_dim(self, inputs, dims=(-1, -2)):
@@ -88,6 +92,16 @@ def test_fp32_mean_dim_unsupported_keepdim_false(self):
             .check_count({"executorch_exir_dialects_edge__ops_aten_mean_dim": 1})
         )
 
+    def test_fp32_mean_dim_unsupported_dtype(self):
+        inputs = (torch.randn(1, 5, 4, 4),)
+        (
+            Tester(self.MeanDim((-1, -2), dtype=torch.float64), inputs)
+            .export()
+            .check_count({"torch.ops.aten.mean.dim": 1})
+            .to_edge_transform_and_lower()
+            .check_count({"executorch_exir_dialects_edge__ops_aten_mean_dim": 1})
+        )
+
     def test_qs8_mean_dim(self):
         inputs = (torch.randn(1, 5, 4, 4),)
         (

backends/xnnpack/utils/utils.py

@@ -60,11 +60,18 @@ def get_input_node(node: torch.fx.Node, input_index: int) -> torch.fx.Node:
 
 def normalize_mean_dims(mean_dims: Sequence[int] | int | None, rank: int) -> List[int]:
     """Return mean dims as non-negative indices for the given rank."""
+    if rank <= 0:
+        raise ValueError(f"Expected rank > 0, got {rank}")
     if mean_dims is None:
         return list(range(rank))
     if isinstance(mean_dims, int):
         mean_dims = [mean_dims]
-    return [dim % rank for dim in mean_dims]
+    normalized_dims = []
+    for dim in mean_dims:
+        if dim < -rank or dim >= rank:
+            raise ValueError(f"Dimension out of range: {dim} for rank {rank}")
+        normalized_dims.append(dim % rank)
+    return normalized_dims
 
 
 def get_relu_fused_node(node: torch.fx.Node) -> Optional[torch.fx.Node]: