Skip to content

[Fix][Relax][Frontend][Torch] Fix torch.round(x, decimals) via from_exported_program and negative-decimals rounding - #20239

Open
siyiweigeHEW wants to merge 1 commit into
apache:mainfrom
siyiweigeHEW:fix/relax-torch-round
Open

[Fix][Relax][Frontend][Torch] Fix torch.round(x, decimals) via from_exported_program and negative-decimals rounding#20239
siyiweigeHEW wants to merge 1 commit into
apache:mainfrom
siyiweigeHEW:fix/relax-torch-round

Conversation

@siyiweigeHEW

@siyiweigeHEW siyiweigeHEW commented Aug 30, 2026

Copy link
Copy Markdown
Contributor

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.roundnearbyint
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.020, 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:

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.

@siyiweigeHEW
siyiweigeHEW force-pushed the fix/relax-torch-round branch from b5f0e3c to c44b474 Compare August 30, 2026 15:52
@siyiweigeHEW siyiweigeHEW changed the title [Relax][Frontend][Torch] Fix torch.round(x, decimals) via from_exported_program and negative-decimals rounding [Fix][Relax][Frontend][Torch] Fix torch.round(x, decimals) via from_exported_program and negative-decimals rounding Aug 30, 2026
@siyiweigeHEW
siyiweigeHEW force-pushed the fix/relax-torch-round branch from c44b474 to b60416b Compare August 30, 2026 21:29

@tlopex tlopex left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.
@tlopex

tlopex commented Sep 3, 2026

Copy link
Copy Markdown
Member

10 ** abs(decimals) is evaluated as a Python integer before it is converted to the target dtype. PyTorch accepts values such as decimals=309, and torch.export emits a valid aten.round.decimals node, but relax.const(10**309, "float32") and the float64 variant raise OverflowError: int too large to convert to float. As a result, these valid exported programs still fail during import. Even larger values can also require constructing an unnecessarily large host integer.

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 decimals value.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

2 participants