Cast mean operand to output dtype before reduce_mean

torch.mean(x, dtype=...) on an integer tensor crashed during conversion
(ValueError: mean[0]: dtype si32 vs f32): replace_mean_default and
replace_mean_dim never honored the dtype kwarg, so reduce_mean ran on the
int operand and produced a result whose type mismatched the node's declared
float output. Cast the operand to the node's output element type before
reducing, mirroring replace_sum_dim_intlist. Adds integer-operand coverage
for aten.mean.default and aten.mean.dim across float32/float16 targets.

Author: arames

coreai_torch/_aten_to_core.py

@@ -2247,6 +2247,9 @@ def replace_mean_default(
 ) -> Value:
     """Computes global mean across all dimensions, returning a scalar tensor."""
     x = _get_operand(values_map, node, 0)
+    target_type = get_output_element_type_from_node(node)
+    if x.type.element_type != target_type:
+        x = coreai.cast(x, target_type)
     all_dims = list(range(x.type.rank))
     return coreai.shrink_dims(coreai.reduce_mean(x, all_dims), all_dims)
 
@@ -2256,6 +2259,9 @@ def replace_mean_dim(
 ) -> Value:
     """Computes mean along specified dimensions."""
     x, axes = _get_operands(values_map, node, [0, 1])
+    target_type = get_output_element_type_from_node(node)
+    if x.type.element_type != target_type:
+        x = coreai.cast(x, target_type)
     keepdim = len(node.args) >= 3 and bool(node.args[2])
     result = coreai.reduce_mean(x, axes)
     return result if keepdim else coreai.shrink_dims(result, axes)

tests/ops/test_ops.py

@@ -2917,6 +2917,70 @@ def forward(self, x: Tensor) -> Tensor:
     await validate_numerical_output(model=model, x=x, dynamic_shapes=dynamic_shapes)
 
 
+@pytest.mark.parametrize("dynamic", [False, True])
+@pytest.mark.parametrize(
+    "x,dim,keepdim,target_dtype",
+    [
+        # mean.default: global mean (dim=None), int32 input → float32 output.
+        # torch promotes the operand to the requested dtype BEFORE averaging.
+        (torch.randint(0, 10, (3, 4), dtype=torch.int32), None, False, torch.float32),
+        (
+            torch.randint(-5, 5, (2, 3, 4), dtype=torch.int32),
+            None,
+            False,
+            torch.float32,
+        ),
+        # mean.dim: reduce along specific dimensions, int32 input → float32.
+        (torch.randint(0, 10, (3, 4), dtype=torch.int32), 1, False, torch.float32),
+        (torch.randint(0, 10, (3, 4), dtype=torch.int32), 0, True, torch.float32),
+        (
+            torch.randint(-5, 5, (2, 3, 4), dtype=torch.int32),
+            [0, 2],
+            False,
+            torch.float32,
+        ),
+        (
+            torch.randint(-5, 5, (2, 3, 4), dtype=torch.int32),
+            [0, 2],
+            True,
+            torch.float32,
+        ),
+        # dtype-agnostic: root cause is independent of the specific int operand
+        # dtype or the specific float target dtype.
+        (torch.randint(0, 10, (3, 4), dtype=torch.int64), 1, False, torch.float32),
+        (torch.randint(0, 10, (3, 4), dtype=torch.int32), None, False, torch.float16),
+    ],
+)
+async def test_mean_dtype_kwarg_integer(
+    x: Tensor,
+    dim: int | list[int] | None,
+    keepdim: bool,
+    target_dtype: torch.dtype,
+    dynamic: bool,
+) -> None:
+    """Regression: torch.mean(x, dtype=...) on an integer tensor must promote the
+    operand to the requested dtype before averaging, not crash at conversion time.
+
+    Covers both aten.mean.default (dim=None) and aten.mean.dim (dim=...) with an
+    explicit float dtype kwarg on an integer operand — the operand must be cast to
+    the output element type before reduce_mean, mirroring sum's dtype handling.
+    Parametrized across int32/int64 operands and float32/float16 targets since the
+    root cause is dtype-agnostic (the cast was simply never attempted)."""
+
+    class MeanDtypeModel(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+
+        def forward(self, x: Tensor) -> Tensor:
+            if dim is None:
+                return torch.mean(x, dtype=target_dtype)
+            return torch.mean(x, dim=dim, keepdim=keepdim, dtype=target_dtype)
+
+    model = MeanDtypeModel().eval()
+    dynamic_shapes = {"x": _all_dims_dynamic(x)} if dynamic else None
+    await validate_numerical_output(model=model, x=x, dynamic_shapes=dynamic_shapes)
+
+
 @pytest.mark.parametrize("dynamic", [False, True])
 @pytest.mark.parametrize("x", [torch.rand(2, 2)])
 @pytest.mark.parametrize("y", [torch.rand(2, 2)])