[Fix][Relax][Frontend][Torch] Fix torch.round(x, decimals) via from_exported_program and negative-decimals rounding - #20239
Conversation
b5f0e3c to
c44b474
Compare
torch.round(x, decimals) via from_exported_program and negative-decimals roundingtorch.round(x, decimals) via from_exported_program and negative-decimals rounding
c44b474 to
b60416b
Compare
tlopex
left a comment
There was a problem hiding this comment.
Please perform the scaling computation in float32 for float16/bfloat16 inputs and cast the result back. Keeping the intermediates in the input dtype causes overflow: for float16 [25, 125], decimals=4 returns inf instead of the original values, while decimals=5 and -5 produce NaNs because the scale becomes inf. Please add low-precision regression tests for both positive and negative decimals.
…ported_program and fix negative-decimals rounding torch.export lowers torch.round(x, decimals) (any explicit decimals, including decimals=0) to aten.round.decimals, but the exported-program convert map only registered round.default. Any explicit decimals made from_exported_program fail with "Unsupported function types ['round.decimals']". Additionally, BaseFXGraphImporter._round scaled every non-zero decimals by round(x * 10**decimals) / 10**decimals. For negative decimals this multiplies by 0.1 / 0.01 / ..., which is inexact in floating point: in float64, torch.round(torch.tensor(25.0), decimals=-1) computed 25 * 0.1 == 2.5000000000000004 and rounded up to 30 instead of 20. Register "round.decimals" in ExportedProgramImporter.create_convert_map and branch the decimals != 0 scale in _round to use an exact integer power of 10: multiply for positive decimals, divide for negative ones (round(x / 10**|d|) * 10**|d|). The ties-to-even inner rounding is already provided by upstream apache#19367 / apache#19368 (tir.round -> nearbyint across backends) and is not changed here. Validated by the verify_patch.py differential harness on the locked build: Part B (fix + ties-to-even inner round, = latest semantics) matches PyTorch for all 28 combinations (from_exported_program / from_fx x decimals {0,1,2,3,-1,-2,-3} x float32/float64), including the previously-failing round(25, -1) == 20 and round(2.25, 1) == 2.2.
b60416b to
8a28308
Compare
|
Please construct or saturate the scale directly in the target floating-point dtype, or handle scale overflow explicitly, and add a regression test with a large |
Fixes: #20231
Summary
torch.exportlowerstorch.round(x, decimals)(any explicitdecimals,including
decimals=0) toaten.round.decimals, while plaintorch.round(x)lowers to
aten.round.default. The Relax Torch frontend registered onlyround.defaultin the exported-program convert map, so any explicitdecimalsmadefrom_exported_programfail outright with"AssertionError: Unsupported function types ['round.decimals']".In addition, the
decimals != 0path inBaseFXGraphImporter._roundalwaysscaled by
round(x * 10**decimals) / 10**decimals. For negativedecimalsthis multiplies by
0.1 / 0.01 / ..., which loses the exact power-of-10 scaleand, in float64, breaks e.g.
torch.round(torch.tensor(25.0, dtype=float64), decimals=-1)(25 * 0.1 == 2.5000000000000004rounds up to30instead of thecorrect
20).This PR registers
round.decimalsin the exported-program convert map and makesthe
decimals != 0scaling use an exact integer power of 10: multiply forpositive decimals, divide for negative ones.
Root cause
from_exported_programrejectstorch.round(x, decimals). Inexported_program_translator.py,ExportedProgramImporter.create_convert_mapmaps
"round.default": self._roundbut noround.decimalsentry. Sincetorch.exportalways emitsaten.round.decimalswhendecimalsis passedexplicitly — even
decimals=0— every such call hits the"Unsupported function types ['round.decimals']"assert in dispatch.Negative decimals round incorrectly.
_roundcomputedscale = relax.const(10**decimals, dtype)and emitteddivide(round(multiply(arg, scale)), scale)for every non-zerodecimals.For
decimals = -1the scale is0.1; multiplying by a non-integerpower of 10 is inexact in floating point, so
torch.round(torch.tensor([25.0], dtype=torch.float64), decimals=-1)produced
30instead of20. (from_fxshares the same_round.)Fix
exported_program_translator.py— add"round.decimals": self._round,right after
"round.default": self._round,inExportedProgramImporter.create_convert_map, so any explicit-decimalstorch.roundconverts through the existing_round.base_fx_graph_translator.py— inBaseFXGraphImporter._round, keep thedecimals == 0fast path, and branch thedecimals != 0scale:decimals > 0:divide(round(multiply(arg, 10**d)), 10**d)(unchanged).decimals < 0:multiply(round(divide(arg, 10**-d)), 10**-d)— divide bythe exact integer power of 10 and multiply back, avoiding the inexact
× 0.1path.Validation
In-tree regression tests (added)
test_round_decimalsintests/python/relax/test_frontend_from_exported_program.py— runs
verify_model_numerically(Relax vs PyTorch) fordecimals 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, 160atdecimals=-1, and2.25 → 2.2).Before the fix,
decimals=0alone fails to import with"Unsupported function types ['round.decimals']".test_round_decimalsintests/python/relax/test_frontend_from_fx.py— samevalues through
from_fx, asserting TVM output matchestorch.roundfor thesame decimals set (this path already dispatched to
_round, but produced thewrong negative-decimals result before the fix).
Differential test
verify_patch.pywas run on the locked pre-#19368 build (rounds half valuesaway from zero). The real fix code is monkey-patched in; the ties-to-even inner
round is reproduced with
te.nearbyintto stand in for latestrelax.op.roundsemantics (#19368):
export(decimals=-1)AssertionError: Unsupported function types ['round.decimals']fx(decimals=-1)[30,130,170]vs torch[20,120,160]Part B matches PyTorch for all 28 combinations — both frontends
(
from_exported_program,from_fx) × 7decimals×float32/float64—including the previously-failing
round(25, -1) == 20andround(2.25, 1) == 2.2cases.Run:
Files changed
python/tvm/relax/frontend/torch/base_fx_graph_translator.py—_round:negative
decimalsnow divide by the exact integer power of 10(
round(x / 10^|d|) * 10^|d|) instead of multiplying by10**decimals(
× 0.1).python/tvm/relax/frontend/torch/exported_program_translator.py— registerround.decimalsin the exported-program convert map.tests/python/relax/test_frontend_from_exported_program.py— addtest_round_decimals.tests/python/relax/test_frontend_from_fx.py— addtest_round_decimals.