From f3c43b9e099478cf3af049de65171a04e030c112 Mon Sep 17 00:00:00 2001
From: Promise Emmanuel Oluwadare
 <95595405+PromiseGameFi@users.noreply.github.com>
Date: Sat, 22 Aug 2026 21:17:34 +0100
Subject: [PATCH] Fix aten mean dtype handling

---
 onnxscript/function_libs/torch_lib/ops/core.py |  7 +++++--
 tests/function_libs/torch_lib/e2e_ops_tests.py | 13 +++++++++++++
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py
index b9f8d1c69e..adcc71645b 100644
--- a/onnxscript/function_libs/torch_lib/ops/core.py
+++ b/onnxscript/function_libs/torch_lib/ops/core.py
@@ -6350,10 +6350,13 @@ def aten_maximum(self: TTensor, other: TTensor) -> TTensor:
     return op.Max(self, other)
 
 
-@torch_op("aten::mean")
-def aten_mean(self: TReal) -> TReal:
+@torch_op("aten::mean", trace_only=True)
+def aten_mean(self: TReal, dtype: int = -1) -> TReal:
     """mean(Tensor self, *, ScalarType? dtype=None) -> Tensor"""
 
+    if dtype != -1:
+        self = op.Cast(self, to=dtype)
+
     result = op.ReduceMean(self)
     return op.Squeeze(result)
 
diff --git a/tests/function_libs/torch_lib/e2e_ops_tests.py b/tests/function_libs/torch_lib/e2e_ops_tests.py
index 1bea819f11..d36bbd65f0 100644
--- a/tests/function_libs/torch_lib/e2e_ops_tests.py
+++ b/tests/function_libs/torch_lib/e2e_ops_tests.py
@@ -1281,6 +1281,19 @@ def forward(self, x):
         )
         _testing.assert_onnx_program(onnx_program)
 
+    def test_mean_dtype_casts_before_reduction(self):
+        class Model(torch.nn.Module):
+            def forward(self, x):
+                return torch.mean(x, dtype=torch.float64)
+
+        model = Model()
+        x = torch.tensor([1e8, 1.0, -1e8], dtype=torch.float32)
+        onnx_program = torch.onnx.export(model, (x,), dynamo=True, optimize=False)
+
+        _testing.assert_onnx_program(onnx_program)
+        actual = onnx_program.call_reference({onnx_program.model.graph.inputs[0].name: x})[0]
+        torch.testing.assert_close(actual, model(x))
+
     def test_aten_histc_float(self):
         class Model(torch.nn.Module):
             def forward(self, x):