[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,55 @@ 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))
+
+        # Build the scale 10**|decimals| directly in `dtype` instead of as a host
+        # Python int. relax.const(10**n, dtype) first materializes 10**n as an
+        # unbounded int, which is both wasteful for large n and, once n >= 309, dies
+        # in the int-to-float conversion with "OverflowError: int too large to
+        # convert to float". PyTorch accepts such decimals (e.g.
+        # torch.round(x, decimals=309)) and exports a valid aten.round.decimals
+        # node, so importing these programs must not crash. Computing the power as a
+        # float and saturating it to inf once it leaves the finite range of `dtype`
+        # is exactly what happens when PyTorch evaluates the same power in the input
+        # dtype.
+        scale_exp = abs(decimals)
+        # Largest exponent for which 10**n is still finite in `dtype`.
+        # (float16/bfloat16 are upcast to float32 above, so dtype is float32 or
+        # float64 here.)
+        max_scale_exp = {"float32": 38, "float64": 308}[dtype]
+        if scale_exp > max_scale_exp:
+            scale = relax.const(float("inf"), dtype)
+        else:
+            scale = relax.const(10.0**scale_exp, dtype)
+
+        if decimals > 0:
+            scaled = relax.op.multiply(arg, scale)
+            rounded = relax.op.round(scaled)
+            result = relax.op.divide(rounded, scale)
+        else:
+            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,98 @@ 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)
+
+
+def test_round_decimals_large():
+    """A large |decimals| must import and run without OverflowError.
+
+    The scale 10**|decimals| used to be built as an unbounded host Python int
+    before being handed to relax.const, whose int-to-float conversion raises
+    OverflowError ("int too large to convert to float") once |decimals| >= 309
+    (10**309 already exceeds the float64 range). PyTorch accepts such decimals and
+    exports a valid aten.round.decimals node, so importing the exported program
+    must not crash on them. The scale is now built directly in the float dtype and
+    saturates to inf once it leaves the finite range, matching PyTorch, whose
+    all-NaN result here comes from the same inf scale.
+    """
+
+    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, 25.0, -0.5, 0.0], dtype=torch.float32)
+    for decimals in (309, -309):
+        exported_program = export(RoundDecimalsModel(decimals).eval(), args=(x,))
+        mod = from_exported_program(exported_program)  # used to raise OverflowError here
+        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()
+
+        # The scale overflows to inf, and IEEE arithmetic turns every element into
+        # NaN in both TVM and PyTorch. Compare the NaN masks and the remaining
+        # (empty here) finite elements separately, since allclose fails on NaN.
+        expected = torch.round(x, decimals=decimals)
+        actual = torch.as_tensor(got)
+        assert torch.equal(torch.isnan(actual), torch.isnan(expected))
+        finite = ~torch.isnan(expected)
+        assert torch.allclose(actual[finite], expected[finite], 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,112 @@ 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_round_decimals_large():
+    """A large |decimals| must import and run without OverflowError.
+
+    The scale 10**|decimals| used to be built as an unbounded host Python int
+    before being handed to relax.const, whose int-to-float conversion raises
+    OverflowError ("int too large to convert to float") once |decimals| >= 309
+    (10**309 already exceeds the float64 range). PyTorch accepts such decimals --
+    torch.round(x, decimals=309) -- and traces a valid round.decimals call, so
+    importing the graph must not crash on them. The scale is now built directly
+    in the float dtype and saturates to inf once it leaves the finite range,
+    matching PyTorch, whose all-NaN result here comes from the same inf scale.
+    """
+    input_info = [([5], "float32")]
+    x = torch.tensor([0.5, 1.5, 25.0, -0.5, 0.0], 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 (309, -309):
+        gm = fx.symbolic_trace(RoundDecimalsModel(decimals).eval())
+        mod = from_fx(gm, input_info)  # used to raise OverflowError here
+        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()
+
+        # The scale overflows to inf, and IEEE arithmetic turns every element into
+        # NaN in both TVM and PyTorch. Compare the NaN masks and the remaining
+        # (empty here) finite elements separately, since allclose fails on NaN.
+        expected = torch.round(x, decimals=decimals)
+        actual = torch.as_tensor(got)
+        assert torch.equal(torch.isnan(actual), torch.isnan(expected))
+        finite = ~torch.isnan(expected)
+        assert torch.allclose(actual[finite], expected[finite], rtol=1e-6, atol=1e-6)
+
+
 def test_size():
     input_info = [([1, 3, 10, 10], "float32")]