From 1e8abd04c57a5b582b2ccd07914c718992c366b3 Mon Sep 17 00:00:00 2001 From: Chen Yufan Date: Thu, 17 Sep 2026 15:58:59 +0800 Subject: [PATCH] [Fix][Relax][Frontend][Torch] Promote the tensor, not the scalar, in binary ops with a Python scalar `_binary_op` built the constant for a Python scalar operand in the tensor's own dtype, so a float scalar against an integer or bool tensor was truncated before the op ran: `x * 0.5` on an int64 tensor became `x * 0`, `x + 0.5` became `x + 0`, `x < 1.5` became `x < 1`, and `bool_tensor * 2.5` stayed bool. torch promotes the other way. A Python scalar takes part in type promotion at a lower priority than a tensor and widens it only when its category is higher: a float scalar promotes an integer or bool tensor to the default float dtype, an int scalar promotes a bool tensor to int64, and otherwise the tensor's dtype wins. Use torch.result_type as the oracle for that rule, cast the tensor when it has to widen, and build the constant in the promoted dtype. Tensor-tensor promotion was already right and is untouched. The two Constant-vs-scalar dispatch branches pre-cast the scalar the same wrong way and now go through the same path. int64 tensor * 0.5 torch float32 [0.5, 1.0, 1.5] before int64 [0, 0, 0] int64 tensor + 0.5 torch float32 [1.5, 2.5, 3.5] before int64 [1, 2, 3] int64 tensor < 1.5 torch [T, F, F] before [F, F, F] bool tensor * 2.5 torch float32 [2.5, 0, 2.5] before bool [T, F, T] bool tensor + 1 torch int64 [2, 1, 2] before InternalError int64 tensor ** 0.5 torch float32 before InternalError Swept 18 binary ops x 8 tensor dtypes x 6 Python scalars (864 programs, 829 that torch accepts), each built with relax.build(llvm) and compared with torch on result dtype and values: before 477 matched 149 wrong dtype or values 203 raised after 645 matched 52 wrong 132 raised no case fails after this change that did not fail before it; 168 repaired Of what is left, 157 are the division family, where true division of two integers has to give a float even for an int scalar and div.Tensor_mode builds its constant with no dtype at all; that is a separate change. The rest are relax rejecting arithmetic on bool tensors, a uint8 tensor against a negative scalar (torch wraps, relax.const raises), and `x ** True` on an integer tensor. test_linspace's expected IR encoded the truncation: torch's decomposition splits the range at `i < 4.5`, which the frontend emitted as `R.less(i, R.const(4, "int64"))`. It now promotes the index to float32 and compares against 4.5, and the expected module is updated to match. Tests: an IR-level check that `int64 + 0.5` emits astype plus a float32 constant, and numeric checks over add/mul/lt/ge/eq (both operand orders) and sub/rsub/pow/ remainder across int64, int32, uint8, float16, float32 and bool tensors with int, float and bool scalars, asserting both the result dtype and the values. 22 of the 44 fail against the previous head; the other 22 are cases where the tensor's dtype wins, and pin that nothing there moved. --- .../torch/base_fx_graph_translator.py | 52 +++++++- .../test_frontend_from_exported_program.py | 121 ++++++++++++++++-- 2 files changed, 159 insertions(+), 14 deletions(-) diff --git a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py index 3ff3b596af3b..7c16a4fe86d0 100644 --- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py +++ b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py @@ -246,6 +246,38 @@ def _emit_torch_reshape(self, x, dims): x = self.block_builder.emit(relax.op.reshape(x, step)) return x + @staticmethod + def _scalar_result_dtype(tensor_dtype, scalar) -> str | None: + """Return the dtype torch gives ``tensor scalar`` for a Python scalar operand. + + A Python scalar takes part in type promotion at a lower priority than a tensor: it + widens the tensor only when it belongs to a higher category. So a float scalar + promotes an integer or bool tensor to the default float dtype, an int scalar + promotes only a bool tensor (to int64), and otherwise the tensor's dtype wins. + Casting the scalar down to the tensor's dtype instead turns ``x * 0.5`` on an + integer tensor into ``x * 0``. Returns None for a dtype torch cannot map. + """ + import torch # type: ignore + + if not isinstance(scalar, bool | int | float): + return None + torch_dtype = { + "float64": torch.float64, + "float32": torch.float32, + "float16": torch.float16, + "bfloat16": torch.bfloat16, + "int64": torch.int64, + "int32": torch.int32, + "int16": torch.int16, + "int8": torch.int8, + "uint8": torch.uint8, + "bool": torch.bool, + }.get(str(tensor_dtype)) + if torch_dtype is None: + return None + promoted = torch.result_type(torch.empty(0, dtype=torch_dtype), scalar) + return str(promoted).replace("torch.", "") + @staticmethod def _promote_common_dtype(lhs_dtype: str | None, rhs_dtype: str | None) -> str | None: """Return the promoted dtype following PyTorch rules, or None if unsupported.""" @@ -707,13 +739,25 @@ def promote_binary_op_args(lhs, rhs): return lhs, rhs elif isinstance(lhs, relax.Expr): assert isinstance(lhs.ty, relax.TensorType) - return lhs, relax.const(rhs, lhs.ty.dtype) + lhs, rhs = promote_scalar_operand(lhs, rhs) + return lhs, rhs elif isinstance(rhs, relax.Expr): assert isinstance(rhs.ty, relax.TensorType) - return relax.const(lhs, rhs.ty.dtype), rhs + rhs, lhs = promote_scalar_operand(rhs, lhs) + return lhs, rhs else: assert False + def promote_scalar_operand(tensor, scalar): + # torch.result_type decides who wins; the tensor is only widened when the + # scalar's category is higher (float scalar vs int tensor, int scalar vs + # bool tensor). The constant is then built in that dtype rather than + # truncated to the tensor's. + target = self._scalar_result_dtype(tensor.ty.dtype, scalar) or tensor.ty.dtype + if str(tensor.ty.dtype) != str(target): + tensor = self.block_builder.emit(relax.op.astype(tensor, target)) + return tensor, relax.const(scalar, target) + def call_binary_op(op, lhs, rhs): lhs, rhs = promote_binary_op_args(lhs, rhs) return self.block_builder.emit(op(lhs, rhs)) @@ -725,9 +769,9 @@ def call_binary_op(op, lhs, rhs): ): return call_binary_op(relax_op, lhs, rhs) elif isinstance(lhs, relax.expr.Constant) and not isinstance(rhs, relax.expr.Constant): - return call_binary_op(relax_op, lhs, relax.const(rhs, dtype=lhs.ty.dtype)) + return call_binary_op(relax_op, lhs, rhs) elif isinstance(rhs, relax.expr.Constant) and not isinstance(lhs, relax.expr.Constant): - return call_binary_op(relax_op, relax.const(lhs, dtype=rhs.ty.dtype), rhs) + return call_binary_op(relax_op, lhs, rhs) return intrinsic_op(lhs, rhs) return convert diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index d78629f5737b..28ef877fd9cb 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/tests/python/relax/test_frontend_from_exported_program.py @@ -1435,6 +1435,103 @@ def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tuple( verify_model(BinaryPromoteRHS(), example_args, {}, expected_promote_rhs) +def test_binary_python_scalar_promotes_the_tensor(): + # A Python scalar only widens the tensor when its category is higher, and then the + # constant has to be built in the promoted dtype. Truncating the scalar to the + # tensor's dtype instead turns ``x + 0.5`` on an int64 tensor into ``x + 0``. + class AddHalf(Module): + def forward(self, x): + return x + 0.5 + + @tvm.script.ir_module + class expected_add_half: + @R.function + def main(x: R.Tensor((3,), dtype="int64")) -> R.Tuple(R.Tensor((3,), dtype="float32")): + with R.dataflow(): + lv: R.Tensor((3,), dtype="float32") = R.astype(x, dtype="float32") + lv1: R.Tensor((3,), dtype="float32") = R.add(lv, R.const(0.5, "float32")) + gv: R.Tuple(R.Tensor((3,), dtype="float32")) = (lv1,) + R.output(gv) + return gv + + verify_model(AddHalf(), (torch.tensor([1, 2, 3]),), {}, expected_add_half) + + +def _scalar_promotion_cases(): + """(tensor dtype, scalar): torch.result_type decides the outcome in every case.""" + return [ + (torch.int64, 0.5), + (torch.int32, 1.5), + (torch.uint8, -0.5), + (torch.int64, 2), + (torch.float16, 2), + (torch.float32, True), + (torch.bool, 2), + (torch.bool, 0.5), + ] + + +def _scalar_input(dtype): + if dtype is torch.bool: + return torch.tensor([True, False, True]) + return torch.tensor([1, 2, 3], dtype=dtype) + + +def _verify_scalar_promotion(model, x): + # Compare both the result dtype and the values: the failure mode this guards is + # ``int * 0.5`` coming back as an int64 tensor of zeros, which a shape check passes. + with torch.no_grad(): + want = model(x) + mod = from_exported_program(export(model, (x,))) + got_dtype = str(mod["main"].ret_ty.fields[0].dtype) + assert got_dtype == str(want.dtype).replace("torch.", ""), ( + f"result dtype {got_dtype}, torch gives {want.dtype}" + ) + verify_model_numerically(model, (x,), rtol=1e-3, atol=1e-3) + + +@pytest.mark.parametrize("op", [operator.add, operator.mul, operator.lt, operator.ge, operator.eq]) +@pytest.mark.parametrize("dtype, scalar", _scalar_promotion_cases()) +def test_binary_python_scalar_promotion_values(op, dtype, scalar): + class Scalar(Module): + def forward(self, x): + return op(x, scalar) + + class ScalarOnTheLeft(Module): + def forward(self, x): + return op(scalar, x) + + x = _scalar_input(dtype) + _verify_scalar_promotion(Scalar(), x) + _verify_scalar_promotion(ScalarOnTheLeft(), x) + + +@pytest.mark.parametrize( + "dtype, scalar", [(torch.int64, 0.5), (torch.int32, 1.5), (torch.int64, 2)] +) +def test_binary_python_scalar_promotion_sub_pow_remainder(dtype, scalar): + # These three reject a bool tensor in torch, so they get their own case list. + class Sub(Module): + def forward(self, x): + return x - scalar + + class RSub(Module): + def forward(self, x): + return scalar - x + + class Pow(Module): + def forward(self, x): + return x**scalar + + class Remainder(Module): + def forward(self, x): + return x % scalar + + x = _scalar_input(dtype) + for model in (Sub(), RSub(), Pow(), Remainder()): + _verify_scalar_promotion(model, x) + + operator_binary_2 = [ (operator.eq, R.equal), (operator.ne, R.not_equal), @@ -8120,16 +8217,20 @@ def main(input: R.Tensor((9, 9), dtype="float32")) -> R.Tuple( lv: R.Tensor((9,), dtype="int64") = R.arange( R.prim_value(0), R.prim_value(9), R.prim_value(1), dtype="int64" ) - lv1: R.Tensor((9,), dtype="bool") = R.less(lv, R.const(4, "int64")) - lv2: R.Tensor((9,), dtype="float32") = R.astype(lv, dtype="float32") - lv3: R.Tensor((9,), dtype="float32") = R.multiply(lv2, R.const(0.125, "float32")) - lv4: R.Tensor((9,), dtype="float32") = R.add(lv3, R.const(0.0, "float32")) - lv5: R.Tensor((9,), dtype="int64") = R.subtract(R.const(8, "int64"), lv) - lv6: R.Tensor((9,), dtype="float32") = R.astype(lv5, dtype="float32") - lv7: R.Tensor((9,), dtype="float32") = R.multiply(lv6, R.const(0.125, "float32")) - lv8: R.Tensor((9,), dtype="float32") = R.subtract(R.const(1.0, "float32"), lv7) - lv9: R.Tensor((9,), dtype="float32") = R.where(lv1, lv4, lv8) - gv: R.Tuple(R.Tensor((9,), dtype="float32")) = (lv9,) + # torch's decomposition splits the range at ``i < 4.5``: the float + # scalar promotes the int64 index to float32 rather than being + # truncated to ``i < 4``. + lv1: R.Tensor((9,), dtype="float32") = R.astype(lv, dtype="float32") + lv2: R.Tensor((9,), dtype="bool") = R.less(lv1, R.const(4.5, "float32")) + lv3: R.Tensor((9,), dtype="float32") = R.astype(lv, dtype="float32") + lv4: R.Tensor((9,), dtype="float32") = R.multiply(lv3, R.const(0.125, "float32")) + lv5: R.Tensor((9,), dtype="float32") = R.add(lv4, R.const(0.0, "float32")) + lv6: R.Tensor((9,), dtype="int64") = R.subtract(R.const(8, "int64"), lv) + lv7: R.Tensor((9,), dtype="float32") = R.astype(lv6, dtype="float32") + lv8: R.Tensor((9,), dtype="float32") = R.multiply(lv7, R.const(0.125, "float32")) + lv9: R.Tensor((9,), dtype="float32") = R.subtract(R.const(1.0, "float32"), lv8) + lv10: R.Tensor((9,), dtype="float32") = R.where(lv2, lv5, lv9) + gv: R.Tuple(R.Tensor((9,), dtype="float32")) = (lv10,) R.output(gv) return gv