Honor the dtype argument of aten::mean when exporting without dim

onnxscript/function_libs/torch_lib/ops/core.py

@@ -6350,21 +6350,33 @@ 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 and dtype is not None:
+        # Cast before reducing so that the accumulation happens in the requested
+        # dtype, matching PyTorch. Casting the result afterwards would keep the
+        # precision loss of the input dtype.
+        self = op.Cast(self, to=dtype)
+
     result = op.ReduceMean(self)
     return op.Squeeze(result)
 
 
 @torch_op("aten::mean", complex=True, trace_only=True)
-def aten_mean_complex(self: TReal) -> TReal:
+def aten_mean_complex(self: TReal, dtype: int = -1) -> TReal:
     """mean(Tensor self, *, ScalarType? dtype=None) -> Tensor"""
 
     rank = len(self.shape) - 1
     dim = op.Constant(value_ints=list(range(rank)))
     result = op.ReduceMean(self, dim, keepdims=False)
+
+    if dtype != -1 and dtype is not None:
+        raise NotImplementedError(
+            "support for the dtype argument is not implemented for complex tensors"
+        )
+
     return result
 
 

tests/function_libs/torch_lib/extra_opinfo.py

@@ -1206,6 +1206,27 @@ def sample_inputs_max_pool3d_with_indices(op_info, device, dtype, requires_grad,
         yield opinfo_core.SampleInput(arg, kwargs=kwargs)
 
 
+def sample_inputs_mean_dtype(op_info, device, dtype, requires_grad, **kwargs):
+    del op_info  # Unused
+    del kwargs  # Unused
+
+    make_arg = functools.partial(
+        torch_testing.make_tensor, device=device, dtype=dtype, requires_grad=requires_grad
+    )
+    for shape in ((S, S), (S,), ()):
+        yield opinfo_core.SampleInput(make_arg(shape), kwargs={"dtype": torch.float64})
+
+    # Precision sensitive values: accumulating in float32 gives 0.0 while accumulating
+    # in float64 gives 1/3, so an implementation that casts only the reduced result
+    # cannot pass this sample.
+    yield opinfo_core.SampleInput(
+        torch.tensor(
+            [[1e8, 1.0, -1e8]], dtype=dtype, device=device, requires_grad=requires_grad
+        ),
+        kwargs={"dtype": torch.float64},
+    )
+
+
 def sample_inputs_native_group_norm(op_info, device, dtype, requires_grad, **kwargs):
     del op_info
     make_arg = functools.partial(
@@ -2746,6 +2767,14 @@ def sample_inputs_masked_scatter(op_info, device, dtype, requires_grad, **kwargs
         sample_inputs_func=sample_inputs_max_pool_empty_strides,
         supports_out=False,
     ),
+    opinfo_core.OpInfo(
+        "ops.aten.mean.dtype",
+        op=torch.ops.aten.mean,
+        aten_name="mean",
+        dtypes=common_dtype.floating_types(),
+        sample_inputs_func=sample_inputs_mean_dtype,
+        supports_out=False,
+    ),
     opinfo_core.OpInfo(
         "ops.aten.native_dropout",
         aten_name="native_dropout",

tests/function_libs/torch_lib/ops_test_data.py

@@ -872,6 +872,7 @@ def _where_input_wrangler(
         matcher=lambda sample: sample.kwargs.get("dim") is not None,
         reason="this Aten overload only accept 1 inputs: self",
     ),
+    TorchLibOpInfo("ops.aten.mean.dtype", core_ops.aten_mean),
     TorchLibOpInfo(
         "mean_dim", core_ops.aten_mean_dim, input_wrangler=_mean_input_wrangler
     ).skip(