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

python/tvm/relax/frontend/torch/base_fx_graph_translator.py

@@ -454,12 +454,37 @@ def _round(self, node: fx.Node) -> relax.Expr:
         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))
+
+        if decimals > 0:
+            scale = relax.const(10**decimals, dtype)
+            scaled = relax.op.multiply(arg, scale)
+            rounded = relax.op.round(scaled)
+            result = relax.op.divide(rounded, scale)
+        else:
+            scale = relax.const(10 ** (-decimals), dtype)
+            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:

python/tvm/relax/frontend/torch/exported_program_translator.py

@@ -1766,6 +1766,7 @@ def create_convert_map(
             "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,

tests/python/relax/test_frontend_from_exported_program.py

@@ -148,6 +148,58 @@ def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
     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)
+
+
 operator_bool_unary = [
     (torch.isinf, R.isinf),
     (torch.isnan, R.isnan),

tests/python/relax/test_frontend_from_fx.py

@@ -2506,6 +2506,71 @@ def main(
     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_size():
     input_info = [([1, 3, 10, 10], "float32")]