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 0943b365ba [Fix][Relax][Frontend][Torch] Fix `torch.round(x,
decimals)` via `from_exported_program` and negative-decimals rounding (#20239)
0943b365ba is described below
commit 0943b365ba0aba8607d5a4ac8c6e2bb913378bce
Author: HuEnwei <[email protected]>
AuthorDate: Fri Sep 11 03:11:12 2026 +0800
[Fix][Relax][Frontend][Torch] Fix `torch.round(x, decimals)` via
`from_exported_program` and negative-decimals rounding (#20239)
Fixes: #20231
## Summary
`torch.export` lowers `torch.round(x, decimals)` (any explicit
`decimals`,
including `decimals=0`) to `aten.round.decimals`, while plain
`torch.round(x)`
lowers to `aten.round.default`. The Relax Torch frontend registered only
`round.default` in the exported-program convert map, so **any** explicit
`decimals` made `from_exported_program` fail outright with
`"AssertionError: Unsupported function types ['round.decimals']"`.
In addition, the `decimals != 0` path in `BaseFXGraphImporter._round`
always
scaled by `round(x * 10**decimals) / 10**decimals`. For **negative**
`decimals`
this multiplies by `0.1 / 0.01 / ...`, which loses the exact power-of-10
scale
and, in float64, breaks e.g. `torch.round(torch.tensor(25.0,
dtype=float64),
decimals=-1)` (`25 * 0.1 == 2.5000000000000004` rounds up to `30`
instead of the
correct `20`).
This PR registers `round.decimals` in the exported-program convert map
and makes
the `decimals != 0` scaling use an exact integer power of 10: multiply
for
positive decimals, **divide** for negative ones.
> Note: the round-half-to-even (ties-to-even) semantics themselves are
already
> provided on latest by upstream #19367 / #19368 (`tir.round` →
`nearbyint`
> across all backends); they are **not** changed by this PR. The fixes
here are
> the `round.decimals` dispatch gap and the negative-decimals scale
precision.
## Root cause
1. **`from_exported_program` rejects `torch.round(x, decimals)`.** In
`exported_program_translator.py`,
`ExportedProgramImporter.create_convert_map`
maps `"round.default": self._round` but no `round.decimals` entry. Since
`torch.export` always emits `aten.round.decimals` when `decimals` is
passed
explicitly — even `decimals=0` — every such call hits the
`"Unsupported function types ['round.decimals']"` assert in dispatch.
2. **Negative decimals round incorrectly.** `_round` computed
`scale = relax.const(10**decimals, dtype)` and emitted
`divide(round(multiply(arg, scale)), scale)` for every non-zero
`decimals`.
For `decimals = -1` the scale is `0.1`; multiplying by a non-integer
power of 10 is inexact in floating point, so
`torch.round(torch.tensor([25.0], dtype=torch.float64), decimals=-1)`
produced `30` instead of `20`. (`from_fx` shares the same `_round`.)
## Fix
- `exported_program_translator.py` — add `"round.decimals":
self._round,`
right after `"round.default": self._round,` in
`ExportedProgramImporter.create_convert_map`, so any explicit-`decimals`
`torch.round` converts through the existing `_round`.
- `base_fx_graph_translator.py` — in `BaseFXGraphImporter._round`, keep
the
`decimals == 0` fast path, and branch the `decimals != 0` scale:
- `decimals > 0`: `divide(round(multiply(arg, 10**d)), 10**d)`
(unchanged).
- `decimals < 0`: `multiply(round(divide(arg, 10**-d)), 10**-d)` —
divide by
the exact integer power of 10 and multiply back, avoiding the inexact
`× 0.1` path.
## Validation
### In-tree regression tests (added)
- `test_round_decimals` in
`tests/python/relax/test_frontend_from_exported_program.py`
— runs `verify_model_numerically` (Relax vs PyTorch) for `decimals in
(0, 1, -1, -2)`
over a value set that exercises ties-to-even half values
(`0.5, 1.5, 2.5, 4.5, -0.5, -2.5`) and the negative-decimals path
(`25.0, 125.0, 165.0` → `20, 120, 160` at `decimals=-1`, and `2.25 →
2.2`).
Before the fix, `decimals=0` alone fails to import with
`"Unsupported function types ['round.decimals']"`.
- `test_round_decimals` in `tests/python/relax/test_frontend_from_fx.py`
— same
values through `from_fx`, asserting TVM output matches `torch.round` for
the
same decimals set (this path already dispatched to `_round`, but
produced the
wrong negative-decimals result before the fix).
### Differential test
`verify_patch.py` was run on the locked pre-#19368 build (rounds half
values
away from zero). The real fix code is monkey-patched in; the
ties-to-even inner
round is reproduced with `te.nearbyint` to stand in for latest
`relax.op.round`
semantics (#19368):
| stage | export + fx × decimals {0,1,2,3,-1,-2,-3} | matched | rejected
| diff elements |
|-------|--------------------------------------------|---------|----------|---------------|
| Part 0 — before fix | `export(decimals=-1)` | import fails:
`AssertionError: Unsupported function types ['round.decimals']` |
| | `fx(decimals=-1)` | `[30,130,170]` vs torch `[20,120,160]` |
| Part A — fix, ties-away inner round (locked build) | 14 | 0 | 22 (all
half-value ties — the #19368 ties-to-even gap, unrelated to this PR) |
| Part B — fix + ties-to-even inner round (= latest) | 28 | 0 | 0 |
Part B matches PyTorch for **all** 28 combinations — both frontends
(`from_exported_program`, `from_fx`) × 7 `decimals` ×
`float32`/`float64` —
including the previously-failing `round(25, -1) == 20` and
`round(2.25, 1) == 2.2` cases.
Run:
```bash
export PATH=/home/shenqingchao/miniconda3/envs/tvm23/bin:$PATH
export
PYTHONPATH=/tmp/tvmffi019:/data/shenqingchao/enwei/familyfuzz/tvm/python
export TVM_LIBRARY_PATH=/data/shenqingchao/enwei/familyfuzz/tvm/build
python results/TVM/deepseek-v4-flash/prove_hum/torch_round/verify_patch.py
```
## Files changed
- `python/tvm/relax/frontend/torch/base_fx_graph_translator.py` —
`_round`:
negative `decimals` now divide by the exact integer power of 10
(`round(x / 10^|d|) * 10^|d|`) instead of multiplying by `10**decimals`
(`× 0.1`).
- `python/tvm/relax/frontend/torch/exported_program_translator.py` —
register
`round.decimals` in the exported-program convert map.
- `tests/python/relax/test_frontend_from_exported_program.py` — add
`test_round_decimals`.
- `tests/python/relax/test_frontend_from_fx.py` — add
`test_round_decimals`.
---
.../frontend/torch/base_fx_graph_translator.py | 57 +++++++++--
.../frontend/torch/exported_program_translator.py | 1 +
.../relax/test_frontend_from_exported_program.py | 92 ++++++++++++++++++
tests/python/relax/test_frontend_from_fx.py | 106 +++++++++++++++++++++
4 files changed, 250 insertions(+), 6 deletions(-)
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 d42590c256..b0682f1f4c 100644
--- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py
+++ b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py
@@ -477,12 +477,57 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
if decimals == 0:
return self.block_builder.emit(relax.op.round(arg))
- # For decimals != 0, use: round(x * 10^decimals) / 10^decimals
- dtype = arg.ty.dtype
- scale = relax.const(10**decimals, dtype)
- scaled = relax.op.multiply(arg, scale)
- rounded = relax.op.round(scaled)
- result = relax.op.divide(rounded, scale)
+ # For decimals != 0, round on the exact power-of-10 scale and scale
back:
+ # round(x * 10^decimals) / 10^decimals. The scaling must always use an
+ # integer power of 10: multiply for positive decimals, divide for
negative
+ # ones. Dividing for negative decimals (instead of multiplying by
+ # 10**decimals, i.e. 0.1 / 0.01 / ...) avoids float precision errors
such as
+ # 25 * 0.1 == 2.5000000000000004 in float64, which would round up to 30
+ # instead of 20 for torch.round(25, -1).
+ #
+ # For float16/bfloat16 inputs the scaling is done in float32 and cast
+ # back, because 10**|decimals| can overflow the input range: 10**4 ==
10000
+ # with 25 * 10000 == 250000 overflows float16 (max 65504) to inf, and
10**5
+ # already overflows float16 to inf, turning decimals=5 and -5 into NaN.
+ input_dtype = arg.ty.dtype
+ dtype = input_dtype
+ if dtype in ("float16", "bfloat16"):
+ dtype = "float32"
+ arg = self.block_builder.emit(relax.op.astype(arg, dtype))
+
+ # Build the scale 10**|decimals| directly in `dtype` instead of as a
host
+ # Python int. relax.const(10**n, dtype) first materializes 10**n as an
+ # unbounded int, which is both wasteful for large n and, once n >=
309, dies
+ # in the int-to-float conversion with "OverflowError: int too large to
+ # convert to float". PyTorch accepts such decimals (e.g.
+ # torch.round(x, decimals=309)) and exports a valid aten.round.decimals
+ # node, so importing these programs must not crash. Computing the
power as a
+ # float and saturating it to inf once it leaves the finite range of
`dtype`
+ # is exactly what happens when PyTorch evaluates the same power in the
input
+ # dtype.
+ scale_exp = abs(decimals)
+ # Largest exponent for which 10**n is still finite in `dtype`.
+ # (float16/bfloat16 are upcast to float32 above, so dtype is float32 or
+ # float64 here.) Note that `dtype` may be a tvm.DataType-like object
whose
+ # repr is "T.float32" instead of a plain str, so select via == rather
than
+ # indexing a str-keyed dict.
+ max_scale_exp = 308 if dtype == "float64" else 38
+ if scale_exp > max_scale_exp:
+ scale = relax.const(float("inf"), dtype)
+ else:
+ scale = relax.const(10.0**scale_exp, dtype)
+
+ if decimals > 0:
+ scaled = relax.op.multiply(arg, scale)
+ rounded = relax.op.round(scaled)
+ result = relax.op.divide(rounded, scale)
+ else:
+ scaled = relax.op.divide(arg, scale)
+ rounded = relax.op.round(scaled)
+ result = relax.op.multiply(rounded, scale)
+
+ if input_dtype in ("float16", "bfloat16"):
+ result = relax.op.astype(result, input_dtype)
return self.block_builder.emit(result)
def _softmax(self, node: fx.Node) -> relax.Var:
diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py
b/python/tvm/relax/frontend/torch/exported_program_translator.py
index e6df019c5c..f01831d63c 100644
--- a/python/tvm/relax/frontend/torch/exported_program_translator.py
+++ b/python/tvm/relax/frontend/torch/exported_program_translator.py
@@ -1771,6 +1771,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
"relu6.default": self._unary_op(relax.op.nn.relu6),
"relu6_.default": self._unary_op(relax.op.nn.relu6),
"round.default": self._round,
+ "round.decimals": self._round,
"rsqrt.default": self._rsqrt,
"scalar_tensor.default": self._scalar_tensor,
"scatter.value": self._scatter_value,
diff --git a/tests/python/relax/test_frontend_from_exported_program.py
b/tests/python/relax/test_frontend_from_exported_program.py
index 75962826f9..565fd836bd 100644
--- a/tests/python/relax/test_frontend_from_exported_program.py
+++ b/tests/python/relax/test_frontend_from_exported_program.py
@@ -148,6 +148,98 @@ def test_basic_unary_ops(pytorch_op, relax_op):
verify_model(UnaryOp(), example_args, {}, expected)
+def test_round_decimals():
+ """torch.round(x, decimals) is exported as aten.round.decimals, which was
missing
+ from the convert map (only round.default was registered) and made any
explicit
+ decimals -- including decimals=0 -- fail with
+ "AssertionError: Unsupported function types ['round.decimals']".
+
+ With the decimals overload registered, torch.round(x, decimals) must
convert and
+ match PyTorch's round-half-to-even results, including negative decimals
+ (round(25, -1) == 20) where the scale-by-0.1 float precision path used to
be wrong.
+ """
+
+ class RoundDecimalsModel(Module):
+ def __init__(self, decimals):
+ super().__init__()
+ self.decimals = decimals
+
+ def forward(self, input):
+ return torch.round(input, decimals=self.decimals)
+
+ # Half values exercise ties-to-even; 25/125/165 exercise the
negative-decimals path.
+ x = torch.tensor(
+ [0.5, 1.5, 2.5, 4.5, -0.5, -2.5, 25.0, 125.0, 165.0, 2.25],
dtype=torch.float32
+ )
+ for decimals in (0, 1, -1, -2):
+ verify_model_numerically(RoundDecimalsModel(decimals).eval(), (x,),
rtol=1e-6, atol=1e-6)
+
+
+def test_round_decimals_low_precision():
+ """Scaling for low-precision inputs must happen in float32 and be cast
back.
+
+ 10**|decimals| can overflow float16: 10**4 == 10000 with 25 * 10000 ==
250000
+ exceeds float16's max of 65504, so scaling in float16 yields inf, and 10**5
+ already overflows float16 (the scale itself becomes inf), turning
decimals=5
+ and -5 into NaN. Upcasting the input to float32 keeps the scaling exact;
the
+ rounded result is cast back to the input dtype.
+ """
+
+ class RoundDecimalsModel(Module):
+ def __init__(self, decimals):
+ super().__init__()
+ self.decimals = decimals
+
+ def forward(self, input):
+ return torch.round(input, decimals=self.decimals)
+
+ x = torch.tensor([0.5, 1.5, 2.5, 2.25, 25.0, 125.0, 165.0, -0.5],
dtype=torch.float16)
+ # Positive decimals exercise the multiply-by-10**d overflow (4, 5);
+ # negative decimals exercise the 10**|d| scale overflowing float16 (-5).
+ for decimals in (2, 4, 5, -2, -4, -5):
+ verify_model_numerically(RoundDecimalsModel(decimals).eval(), (x,),
rtol=1e-6, atol=1e-6)
+
+
+def test_round_decimals_large():
+ """A large |decimals| must import and run without OverflowError.
+
+ The scale 10**|decimals| used to be built as an unbounded host Python int
+ before being handed to relax.const, whose int-to-float conversion raises
+ OverflowError ("int too large to convert to float") once |decimals| >= 309
+ (10**309 already exceeds the float64 range). PyTorch accepts such decimals
and
+ exports a valid aten.round.decimals node, so importing the exported program
+ must not crash on them. The scale is now built directly in the float dtype
and
+ saturates to inf once it leaves the finite range, matching PyTorch, whose
+ all-NaN result here comes from the same inf scale.
+ """
+
+ class RoundDecimalsModel(Module):
+ def __init__(self, decimals):
+ super().__init__()
+ self.decimals = decimals
+
+ def forward(self, input):
+ return torch.round(input, decimals=self.decimals)
+
+ x = torch.tensor([0.5, 1.5, 25.0, -0.5, 0.0], dtype=torch.float32)
+ for decimals in (309, -309):
+ exported_program = export(RoundDecimalsModel(decimals).eval(),
args=(x,))
+ mod = from_exported_program(exported_program) # used to raise
OverflowError here
+ ex = relax.build(mod, target="llvm")
+ vm = relax.VirtualMachine(ex, tvm.cpu())
+ tvm_out = vm["main"](tvm.runtime.tensor(x.numpy()))
+ got = tvm_out.numpy() if hasattr(tvm_out, "numpy") else
tvm_out[0].numpy()
+
+ # The scale overflows to inf, and IEEE arithmetic turns every element
into
+ # NaN in both TVM and PyTorch. Compare the NaN masks and the remaining
+ # (empty here) finite elements separately, since allclose fails on NaN.
+ expected = torch.round(x, decimals=decimals)
+ actual = torch.as_tensor(got)
+ assert torch.equal(torch.isnan(actual), torch.isnan(expected))
+ finite = ~torch.isnan(expected)
+ assert torch.allclose(actual[finite], expected[finite], rtol=1e-6,
atol=1e-6)
+
+
operator_bool_unary = [
(torch.isinf, R.isinf),
(torch.isnan, R.isnan),
diff --git a/tests/python/relax/test_frontend_from_fx.py
b/tests/python/relax/test_frontend_from_fx.py
index 7189b3cb24..a0ba7971db 100644
--- a/tests/python/relax/test_frontend_from_fx.py
+++ b/tests/python/relax/test_frontend_from_fx.py
@@ -2556,6 +2556,112 @@ def test_div_mode():
verify_model(DivFloorModel(), input_info, {}, expected_div_floor)
+def test_round_decimals():
+ """torch.round(x, decimals) through from_fx must match PyTorch's
round-half-to-even
+ results, including negative decimals (round(25, -1) == 20). The previous
+ scale-by-10**decimals implementation multiplied by 0.1 for negative
decimals, which
+ is numerically wrong: 25 * 0.1 == 2.5000000000000004 in float64 rounds up
to 30.
+ """
+ input_info = [([10], "float32")]
+ x = torch.tensor(
+ [0.5, 1.5, 2.5, 4.5, -0.5, -2.5, 25.0, 125.0, 165.0, 2.25],
dtype=torch.float32
+ )
+
+ class RoundDecimalsModel(Module):
+ def __init__(self, decimals):
+ super().__init__()
+ self.decimals = decimals
+
+ def forward(self, input):
+ return torch.round(input, decimals=self.decimals)
+
+ for decimals in (0, 1, -1, -2):
+ gm = fx.symbolic_trace(RoundDecimalsModel(decimals).eval())
+ mod = from_fx(gm, input_info)
+ ex = relax.build(mod, target="llvm")
+ vm = relax.VirtualMachine(ex, tvm.cpu())
+ tvm_out = vm["main"](tvm.runtime.tensor(x.numpy()))
+ got = tvm_out.numpy() if hasattr(tvm_out, "numpy") else
tvm_out[0].numpy()
+ tvm.testing.assert_allclose(
+ got, torch.round(x, decimals=decimals).numpy(), rtol=1e-6,
atol=1e-6
+ )
+
+
+def test_round_decimals_low_precision():
+ """Scaling for low-precision inputs must happen in float32 and be cast
back.
+
+ 10**|decimals| can overflow float16: 10**4 == 10000 with 25 * 10000 ==
250000
+ exceeds float16's max of 65504, so scaling in float16 yields inf, and 10**5
+ already overflows float16 (the scale itself becomes inf), turning
decimals=5
+ and -5 into NaN. Upcasting the input to float32 keeps the scaling exact;
the
+ rounded result is cast back to the input dtype.
+ """
+ input_info = [([8], "float16")]
+ x = torch.tensor([0.5, 1.5, 2.5, 2.25, 25.0, 125.0, 165.0, -0.5],
dtype=torch.float16)
+
+ class RoundDecimalsModel(Module):
+ def __init__(self, decimals):
+ super().__init__()
+ self.decimals = decimals
+
+ def forward(self, input):
+ return torch.round(input, decimals=self.decimals)
+
+ # Positive decimals exercise the multiply-by-10**d overflow (4, 5);
+ # negative decimals exercise the 10**|d| scale overflowing float16 (-5).
+ for decimals in (2, 4, 5, -2, -4, -5):
+ gm = fx.symbolic_trace(RoundDecimalsModel(decimals).eval())
+ mod = from_fx(gm, input_info)
+ ex = relax.build(mod, target="llvm")
+ vm = relax.VirtualMachine(ex, tvm.cpu())
+ tvm_out = vm["main"](tvm.runtime.tensor(x.numpy()))
+ got = tvm_out.numpy() if hasattr(tvm_out, "numpy") else
tvm_out[0].numpy()
+ tvm.testing.assert_allclose(
+ got, torch.round(x, decimals=decimals).numpy(), rtol=1e-6,
atol=1e-6
+ )
+
+
+def test_round_decimals_large():
+ """A large |decimals| must import and run without OverflowError.
+
+ The scale 10**|decimals| used to be built as an unbounded host Python int
+ before being handed to relax.const, whose int-to-float conversion raises
+ OverflowError ("int too large to convert to float") once |decimals| >= 309
+ (10**309 already exceeds the float64 range). PyTorch accepts such decimals
--
+ torch.round(x, decimals=309) -- and traces a valid round.decimals call, so
+ importing the graph must not crash on them. The scale is now built directly
+ in the float dtype and saturates to inf once it leaves the finite range,
+ matching PyTorch, whose all-NaN result here comes from the same inf scale.
+ """
+ input_info = [([5], "float32")]
+ x = torch.tensor([0.5, 1.5, 25.0, -0.5, 0.0], dtype=torch.float32)
+
+ class RoundDecimalsModel(Module):
+ def __init__(self, decimals):
+ super().__init__()
+ self.decimals = decimals
+
+ def forward(self, input):
+ return torch.round(input, decimals=self.decimals)
+
+ for decimals in (309, -309):
+ gm = fx.symbolic_trace(RoundDecimalsModel(decimals).eval())
+ mod = from_fx(gm, input_info) # used to raise OverflowError here
+ ex = relax.build(mod, target="llvm")
+ vm = relax.VirtualMachine(ex, tvm.cpu())
+ tvm_out = vm["main"](tvm.runtime.tensor(x.numpy()))
+ got = tvm_out.numpy() if hasattr(tvm_out, "numpy") else
tvm_out[0].numpy()
+
+ # The scale overflows to inf, and IEEE arithmetic turns every element
into
+ # NaN in both TVM and PyTorch. Compare the NaN masks and the remaining
+ # (empty here) finite elements separately, since allclose fails on NaN.
+ expected = torch.round(x, decimals=decimals)
+ actual = torch.as_tensor(got)
+ assert torch.equal(torch.isnan(actual), torch.isnan(expected))
+ finite = ~torch.isnan(expected)
+ assert torch.allclose(actual[finite], expected[finite], rtol=1e-6,
atol=1e-6)
+
+
def test_size():
input_info = [([1, 3, 10, 10], "float32")]
