This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new bc8947b7dd [Fix][Relax][Frontend][Torch] Honor the `dtype` argument of
`aten.mean` (`torch.Tensor.mean` / `torch.mean`) (#20241)
bc8947b7dd is described below
commit bc8947b7ddd02f14680e0b3d6ea3ea7cb93f3acd
Author: HuEnwei <[email protected]>
AuthorDate: Wed Sep 2 09:52:55 2026 +0800
[Fix][Relax][Frontend][Torch] Honor the `dtype` argument of `aten.mean`
(`torch.Tensor.mean` / `torch.mean`) (#20241)
Fixes: #20230
## Summary
PyTorch's `Tensor.mean` (and the equivalent `torch.mean`) accepts an
optional keyword-only `dtype` argument that controls **both** the
accumulation type and the output type. `torch.export` preserves it on
the
`aten.mean.dim` / `aten.mean.default` nodes, but the `_mean` converter
in
`base_fx_graph_translator.py` never reads `node.kwargs["dtype"]`, so the
argument is silently dropped and the output keeps the input dtype. For
example, `x.mean(dim=1, dtype=torch.float64)` on an fp32 `x` returns
`float32` instead of `float64`, and in the fp16→fp32 direction the
values
are also accumulated at the wrong (lower) precision.
This PR makes `_mean` honor `dtype` by casting the input to the
requested
dtype before reducing — the same pattern `_sum` already uses.
## Root cause
`_mean`
(`python/tvm/relax/frontend/torch/base_fx_graph_translator.py:1646`)
only reads `dim` and `keepdim`:
```python
def _mean(self, node: fx.Node) -> relax.Var:
args = self.retrieve_args(node)
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)
return self.block_builder.emit(relax.op.mean(x, dim, keepdims=keepdim))
```
`relax.op.mean` (`python/tvm/relax/op/statistical.py:54`) has no
`out_dtype` parameter, so the requested dtype cannot be propagated.
`torch.export` does preserve the `dtype` kwarg on the node:
```
%mean : call_function[target=torch.ops.aten.mean.dim](args = (%x, [1]),
kwargs = {dtype: torch.float64})
```
The same converter is registered for `mean.dim` / `mean.default`
(exported-program path) and for `mean` (fx-trace path), so both entry
points are affected.
## Fix
When `dtype` is present on the node, cast the input to the requested
dtype
before emitting the mean, mirroring the existing `_sum` handling:
```python
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))
```
Casting the input and then reducing is exactly PyTorch's documented
semantics: the `dtype` argument controls the accumulation type (the
input
is converted to `dtype` before the reduction), so no precision is lost
between the cast and the mean.
## Validation
### In-tree regression test
Extended `test_mean` in `tests/python/relax/test_frontend_from_fx.py`
with
`MeanDtype`, an fx-traced module calling `input.mean(-1,
dtype=torch.float64)`, and its expected IR: an `R.astype` to `"float64"`
followed by `R.mean` producing a `float64` output. The produced module
is
checked with `tvm.ir.assert_structural_equal` via the file's
`verify_model`
helper.
### Differential test
The prove_hum differential harness (`1复现_torch_mean.py`, converting
`torch.export` output with `from_exported_program` and comparing against
native PyTorch 2.10.0) was run on a runnable build of this change — 15
cases:
- **Baseline (7)**: global / single-dim / multi-dim / keepdim /
negative-axis `mean` without `dtype` — unchanged, all match native
PyTorch (dtype, shape, values).
- **Same-dtype control (1)**: `mean(dtype=fp32)` on fp32 — matches.
- **`dtype != input` (6)**: `fp32→fp64`, `fp64→fp32`, `fp16→fp32`,
`fp32→fp16`, including global and multi-dim + keepdim variants.
All 6 `dtype != input` cases now produce the requested output dtype in
every direction (pre-fix, all 6 returned the input dtype). Values match
to
native PyTorch within fp16/fp32 rounding precision. Both converter entry
paths — `from_exported_program` and `from_fx` — were verified.
## Files changed
- `python/tvm/relax/frontend/torch/base_fx_graph_translator.py` —
`_mean`
reads `node.kwargs["dtype"]` and casts the input with
`relax.op.astype` before emitting `relax.op.mean`.
- `tests/python/relax/test_frontend_from_fx.py` — add `MeanDtype` /
`ExpectedDtype` coverage to `test_mean`.
---
.../tvm/relax/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 2afbd009e7..f5f8afd441 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 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
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 1d086cf043..0759aab1ee 100644
--- a/tests/python/relax/test_frontend_from_fx.py
+++ b/tests/python/relax/test_frontend_from_fx.py
@@ -5070,8 +5070,24 @@ def test_mean():
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():
