From e8762635f6554aa483801fb7fec6b17487bb20a3 Mon Sep 17 00:00:00 2001 From: HuEnwei Date: Sun, 30 Aug 2026 23:14:51 +0800 Subject: [PATCH] [Relax][Frontend][Torch] Honor dtype argument in aten.mean lowering PyTorch's `Tensor.mean` accepts an optional keyword-only `dtype` argument that controls both the accumulation and the output type. `torch.export` preserves it on the `aten.mean.dim` / `aten.mean.default` nodes, but the `_mean` converter never reads `node.kwargs["dtype"]`, so the argument is silently dropped and the output keeps the input dtype (e.g. fp32 mean(dtype=fp64) returns fp32). Match the existing `_sum` handling: when `dtype` is given, cast the input to the requested dtype with `relax.op.astype` before reducing. `relax.op.mean` has no `out_dtype`, so the cast-then-reduce form also matches PyTorch's documented semantics (cast input, then accumulate). --- .../frontend/torch/base_fx_graph_translator.py | 5 +++++ tests/python/relax/test_frontend_from_fx.py | 16 ++++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py index d600987cdd7b..d4b33c095c44 100644 --- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py +++ b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py @@ -1648,6 +1648,11 @@ def _mean(self, node: fx.Node) -> relax.Var: x = args[0] dim = args[1] if len(node.args) > 1 else node.kwargs.get("dim", None) keepdim = args[2] if len(node.args) > 2 else node.kwargs.get("keepdim", False) + dtype = node.kwargs.get("dtype", None) + if dtype is not None: + x = self.block_builder.emit( + relax.op.astype(x, self._convert_data_type(dtype, self.env)) + ) return self.block_builder.emit(relax.op.mean(x, dim, keepdims=keepdim)) def _median(self, node: fx.Node) -> relax.Var: diff --git a/tests/python/relax/test_frontend_from_fx.py b/tests/python/relax/test_frontend_from_fx.py index a489977958c7..03a091735a11 100644 --- a/tests/python/relax/test_frontend_from_fx.py +++ b/tests/python/relax/test_frontend_from_fx.py @@ -5070,8 +5070,24 @@ def main(inp_0: R.Tensor((256, 256), dtype="float32")) -> R.Tensor( R.output(gv) return gv + class MeanDtype(Module): + def forward(self, input): + return input.mean(-1, dtype=torch.float64) + + @I.ir_module + class ExpectedDtype: + @R.function + def main(inp_0: R.Tensor((256, 256), dtype="float32")) -> R.Tensor((256,), dtype="float64"): + with R.dataflow(): + lv: R.Tensor((256, 256), dtype="float64") = R.astype(inp_0, dtype="float64") + lv1: R.Tensor((256,), dtype="float64") = R.mean(lv, axis=[-1], keepdims=False) + gv: R.Tensor((256,), dtype="float64") = lv1 + R.output(gv) + return gv + verify_model(Mean(), [([256, 256], "float32")], {}, Expected1) verify_model(MeanKeepDim(), [([256, 256], "float32")], {}, Expected2) + verify_model(MeanDtype(), [([256, 256], "float32")], {}, ExpectedDtype) def test_cat():