diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index b9f8d1c69e..81f71d53e7 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/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 diff --git a/tests/function_libs/torch_lib/extra_opinfo.py b/tests/function_libs/torch_lib/extra_opinfo.py index 188606bdd6..cea7f03b51 100644 --- a/tests/function_libs/torch_lib/extra_opinfo.py +++ b/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", diff --git a/tests/function_libs/torch_lib/ops_test_data.py b/tests/function_libs/torch_lib/ops_test_data.py index 62ad59c446..3658a35b9d 100644 --- a/tests/function_libs/torch_lib/ops_test_data.py +++ b/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(