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

[siyiweigeHEW]

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:

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.

[tlopex]

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.

[tlopex]

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.

[siyiweigeHEW]

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.

Thanks for catching these edge cases.

I agree that the current implementation should not perform the scaling intermediates in the input dtype for float16/bfloat16, since large positive or negative decimals can cause the scale or intermediate result to overflow. I’ll update the implementation to perform the scaling computation in float32 for low-precision inputs and cast the final result back to the original dtype. I’ll also add regression tests covering both positive and negative decimals for float16 and bfloat16.

For very large decimals such as 309, I’ll also avoid constructing 10**decimals as a large Python integer before converting it to the target dtype. I’ll handle scale construction/overflow explicitly in the target floating-point dtype so that valid torch.round inputs do not fail during import.

I’ll push an updated commit with these changes and the corresponding regression tests. Thanks again for pointing these out.