[Fix][Relax][Frontend][Torch] Honor the dtype argument of aten.mean (torch.Tensor.mean / torch.mean)

[siyiweigeHEW]

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:

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:

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.