siyiweigeHEW opened a new issue, #20230:
URL: https://github.com/apache/tvm/issues/20230
### Expected behavior
`torch.Tensor.mean` (and the equivalent `torch.mean`) accepts an optional
keyword-only `dtype` argument that controls both the accumulation type and the
output type:
```python
>>> x = torch.randn(2, 3, 4, dtype=torch.float32)
>>> x.mean(dim=1, dtype=torch.float64).dtype
torch.float64
```
When a model uses `mean(..., dtype=...)` with `dtype` different from the
input dtype, the TVM PyTorch frontend should produce a tensor of the requested
dtype, matching native PyTorch.
### Actual behavior
`tvm.relax.frontend.torch.from_exported_program` silently drops the `dtype`
argument. The model imports, builds, and runs without any error, but the output
tensor has the **input** dtype instead of the requested dtype. In the `fp16 →
fp32` case the values are also computed at the wrong (lower) precision.
The converter `_mean`
(`python/tvm/relax/frontend/torch/base_fx_graph_translator.py:1646`) only reads
`dim` and `keepdim` and never reads `node.kwargs["dtype"]`:
```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 either, so there is no way the dtype is honored. The same converter
is registered for `mean.dim` / `mean.default`
(`exported_program_translator.py:1915-1916`) and for `mean`
(`fx_translator.py:968`).
`torch.export` does preserve the `dtype` argument on the `aten.mean` node,
so this is reachable from ordinary user code:
```
%mean : call_function[target=torch.ops.aten.mean.dim](args = (%x, [1]),
kwargs = {dtype: torch.float64})
```
### Environment
- OS: Linux
- TVM: v0.24.dev0 (main branch, commit `262c6d2e0`; bug also present in the
latest source at `390af87345`)
- Python: 3.11
- torch: 2.10.0+cu128
### Steps to reproduce
```python
"""Repro: torch.mean(dtype=...) output dtype dropped by TVM relax PyTorch
frontend."""
import torch
import torch.nn as nn
import numpy as np
import tvm
from tvm import relax
from tvm.relax.frontend.torch import from_exported_program
class M(nn.Module):
def forward(self, x):
return x.mean(dim=1, dtype=torch.float64)
m = M().eval()
x = torch.randn(2, 3, 4, dtype=torch.float32)
with torch.no_grad():
ref = m(x) # float64, shape (2, 4)
print("torch ref :", ref.dtype, tuple(ref.shape))
exp = torch.export.export(m.cpu(), (x.cpu(),))
mod = from_exported_program(exp)
ex = relax.build(mod, target="llvm")
vm = relax.VirtualMachine(ex, tvm.cpu())
out = vm["main"](x.numpy())[0].numpy()
print("tvm out :", out.dtype, out.shape) # float32, shape (2, 4)
<-- wrong dtype
```
Actual output:
```
torch ref : torch.float64 (2, 4)
tvm out : float32 (2, 4)
```
The full differential matrix (`prove_hum/torch_mean/1复现_torch_mean.py`)
shows 8/8 baseline cases (global / single / multi-dim / keepdim / negative axis
/ same-dtype control) matching native PyTorch exactly, and 6/6 `dtype != input
dtype` cases diverging in output dtype:
| input | call | torch output | TVM output |
|---|---|---|---|
| fp32 | `mean(dim=1, dtype=fp64)` | `float64` | `float32` |
| fp32 | `mean(dtype=fp64)` | `float64` | `float32` |
| fp32 | `mean(dim=(1,2), keepdim=True, dtype=fp64)` | `float64` | `float32`
|
| fp64 | `mean(dim=1, dtype=fp32)` | `float32` | `float64` |
| fp16 | `mean(dim=1, dtype=fp32)` | `float32` | `float16` |
| fp32 | `mean(dim=1, dtype=fp16)` | `float16` | `float32` |
### Triage
* needs-triage
* bug
* relax
* frontend/torch
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
