siyiweigeHEW opened a new pull request, #20239:
URL: https://github.com/apache/tvm/pull/20239
# [Relax][Frontend][Torch] Fix `torch.round(x, decimals)` via
`from_exported_program` and negative-decimals rounding
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`.
--
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]
