From 1e8abd04c57a5b582b2ccd07914c718992c366b3 Mon Sep 17 00:00:00 2001 From: Chen Yufan Date: Thu, 17 Sep 2026 15:58:59 +0800 Subject: [PATCH 1/2] [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 From ed4b454e6ae2bfc12c580da4b8a7f903f161d50b Mon Sep 17 00:00:00 2001 From: Chen Yufan Date: Thu, 17 Sep 2026 16:16:27 +0800 Subject: [PATCH 2/2] [Fix][Relax][Frontend][Torch] Follow torch's dtype rules for the division family Three converters disagreed with torch on what a division returns: - `div.Tensor` / `div.Scalar` went through the generic `_binary_op`, which keeps the promoted dtype. torch's true division always yields a floating result, so `int64 / 2` and `int64 / int64` are float32 there and were an integer quotient here: `[3, 4, 5] / 2` came back as int64 `[1, 2, 2]` instead of float32 `[1.5, 2, 2.5]`. - `div.Tensor_mode` (which `x // 2` and `torch.div(x, 2, rounding_mode=...)` decompose to) built its scalar constant with `relax.const(inp_2)` and no dtype, i.e. int32, so every float tensor and every non-int32 integer tensor failed the same-dtype check: `x // 2` raised `TypeError` for float32 and int64 inputs alike. The int32 case passed only because relax.const's default dtype happens to be int32. - `reciprocal.default`, which `scalar / x` decomposes to, divided `const(1, x.dtype)` by x, so `2 / int_tensor` was an integer quotient as well. The converter was duplicated in both translators; there is now one in the base class. The two promotion closures inside `_binary_op` become methods (`_promote_binary_operands`, `_promote_scalar_operand`) so the division converters share them, and `_true_division_operands` adds the one rule true division has on top of `torch.result_type`: an integral or bool pair is cast to the default float dtype. `div.Tensor` / `div.Scalar` dispatch to a new `_true_divide`; `_div` promotes its operands the same way and then keeps the promoted dtype for `floor` and `trunc`. Integer division in relax truncates toward zero, so `trunc` on an integer pair is a plain divide; `floor` is `floor_divide`; floats go through divide + trunc as before. int64 [3, 4, 5] / 2 torch float32 [1.5, 2.0, 2.5] before int64 [1, 2, 2] bool [T, T, F] / 2 torch float32 [0.5, 0.5, 0.0] before bool 2 / int64 [3, 4, 5] torch float32 [0.67, 0.5, 0.4] before int64 [0, 0, 0] int64 [-7, -3, 3, 7] // 2 torch int64 [-4, -2, 1, 3] before TypeError float32 x // 2 torch float32 before TypeError torch.div(x, 2, "trunc") torch int64 [-3, -1, 1, 3] before TypeError Same 864-program sweep as the previous commit (18 binary ops x 8 dtypes x 6 Python scalars, built with relax.build(llvm), result dtype and values compared with torch), measured against that commit as the base: base 645 matched 52 wrong dtype or values 132 raised after 799 matched 0 wrong 30 raised no case fails after this change that did not fail before it; 154 repaired The 30 left are the same pre-existing edges as before: relax rejecting arithmetic on a bool tensor with a bool scalar, a uint8 tensor against a negative scalar (torch wraps, relax.const raises), and `x ** True` on an integer tensor. Tests: an IR-level check that `int64 / 2` casts both operands to float32 before the divide; numeric checks of `x / s`, `s / x`, `x / (x + 1)` and `torch.reciprocal(x)` over int64, int32, uint8 and bool tensors; and `x // 2`, `torch.div(..., "floor")`, `torch.div(..., "trunc")` and a tensor divisor on `[-7, -3, 3, 7]` for int64, int32 and float32, where the negative inputs separate floor from trunc. 8 of 9 fail against the previous head; the int32 rounding-mode case passes there for the int32-default reason above and pins that it keeps working. --- .../torch/base_fx_graph_translator.py | 139 +++++++++++------- .../torch/exported_program_translator.py | 8 +- .../tvm/relax/frontend/torch/fx_translator.py | 6 +- .../test_frontend_from_exported_program.py | 74 ++++++++++ 4 files changed, 161 insertions(+), 66 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 7c16a4fe86d0..9c49abb913b4 100644 --- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py +++ b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py @@ -656,6 +656,13 @@ def _round(self, node: fx.Node) -> relax.Expr: result = relax.op.astype(result, input_dtype) return self.block_builder.emit(result) + def _reciprocal(self, node: fx.Node) -> relax.Var: + # torch.reciprocal is 1 / x under true-division rules: an integer or bool input + # comes back as the default float dtype, not as an integer quotient. + x = self.env[node.args[0]] + one, x = self._true_division_operands(1, x) + return self.block_builder.emit(relax.op.divide(one, x)) + def _softmax(self, node: fx.Node) -> relax.Var: x = self.env[node.args[0]] dim = node.args[1] if len(node.args) > 1 else node.kwargs.get("dim", -1) @@ -719,47 +726,63 @@ def convert(node: fx.Node) -> relax.Var: ########## Binary Ops ########## + def _promote_scalar_operand(self, tensor, scalar): + """Widen ``tensor`` if a Python ``scalar`` outranks it, and build the constant. + + 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 _promote_binary_operands(self, lhs, rhs): + """Bring the two operands of a binary op to torch's promoted dtype.""" + if isinstance(lhs, relax.Expr) and isinstance(rhs, relax.Expr): + lhs_si = getattr(lhs, "ty", None) + rhs_si = getattr(rhs, "ty", None) + if isinstance(lhs_si, relax.TensorType) and isinstance(rhs_si, relax.TensorType): + target_dtype = self._promote_common_dtype(lhs_si.dtype, rhs_si.dtype) + if target_dtype is not None: + if lhs_si.dtype != target_dtype: + lhs = self.block_builder.emit(relax.op.astype(lhs, target_dtype)) + if rhs_si.dtype != target_dtype: + rhs = self.block_builder.emit(relax.op.astype(rhs, target_dtype)) + return lhs, rhs + elif isinstance(lhs, relax.Expr): + assert isinstance(lhs.ty, relax.TensorType) + return self._promote_scalar_operand(lhs, rhs) + elif isinstance(rhs, relax.Expr): + assert isinstance(rhs.ty, relax.TensorType) + rhs, lhs = self._promote_scalar_operand(rhs, lhs) + return lhs, rhs + else: + assert False + + def _true_division_operands(self, lhs, rhs): + """Promote for ``a / b``: integer and bool operands divide as the default float. + + torch's true division always produces a floating result -- ``int64 / int64`` and + ``int64 / 2`` are float32, not a truncating integer quotient -- so after the usual + promotion an integral or bool pair is cast to the default float dtype. + """ + lhs, rhs = self._promote_binary_operands(lhs, rhs) + dtype = getattr(getattr(lhs, "ty", None), "dtype", None) + if dtype is not None and ( + dtype.matches_code(DataTypeCode.INT, DataTypeCode.UINT) or str(dtype) == "bool" + ): + lhs = self.block_builder.emit(relax.op.astype(lhs, "float32")) + rhs = self.block_builder.emit(relax.op.astype(rhs, "float32")) + return lhs, rhs + def _binary_op(self, relax_op: Callable, intrinsic_op: Callable) -> Callable: from torch import fx def convert(node: fx.Node) -> relax.Var: - def promote_binary_op_args(lhs, rhs): - if isinstance(lhs, relax.Expr) and isinstance(rhs, relax.Expr): - lhs_si = getattr(lhs, "ty", None) - rhs_si = getattr(rhs, "ty", None) - if isinstance(lhs_si, relax.TensorType) and isinstance( - rhs_si, relax.TensorType - ): - target_dtype = self._promote_common_dtype(lhs_si.dtype, rhs_si.dtype) - if target_dtype is not None: - if lhs_si.dtype != target_dtype: - lhs = self.block_builder.emit(relax.op.astype(lhs, target_dtype)) - if rhs_si.dtype != target_dtype: - rhs = self.block_builder.emit(relax.op.astype(rhs, target_dtype)) - return lhs, rhs - elif isinstance(lhs, relax.Expr): - assert isinstance(lhs.ty, relax.TensorType) - lhs, rhs = promote_scalar_operand(lhs, rhs) - return lhs, rhs - elif isinstance(rhs, relax.Expr): - assert isinstance(rhs.ty, relax.TensorType) - 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) + lhs, rhs = self._promote_binary_operands(lhs, rhs) return self.block_builder.emit(op(lhs, rhs)) lhs, rhs = self.retrieve_args(node) @@ -797,31 +820,37 @@ def _pow(self, node: fx.Node) -> relax.Var: return result return self._binary_op(relax.op.power, operator.pow)(node) + def _true_divide(self, node: fx.Node) -> relax.Var: + lhs, rhs = self.retrieve_args(node) + if not isinstance(lhs, relax.Expr) and not isinstance(rhs, relax.Expr): + return operator.truediv(lhs, rhs) + lhs, rhs = self._true_division_operands(lhs, rhs) + return self.block_builder.emit(relax.op.divide(lhs, rhs)) + def _div(self, node: fx.Node) -> relax.Var: args = self.retrieve_args(node) - inp_1 = args[0] - inp_2 = args[1] - - # Handle scalar cases - if isinstance(inp_2, int | float): - inp_2 = relax.const(inp_2) - - # Get rounding_mode from node kwargs + lhs, rhs = args[0], args[1] rounding_mode = args[2] if len(node.args) > 2 else node.kwargs.get("rounding_mode", None) - # Perform division based on rounding mode if rounding_mode is None: - # True division (normal float division) - return self.block_builder.emit(relax.op.divide(inp_1, inp_2)) - elif rounding_mode == "floor": - # Floor division - return self.block_builder.emit(relax.op.floor_divide(inp_1, inp_2)) - elif rounding_mode == "trunc": - # Trunc division: perform true division then truncate - true_div = self.block_builder.emit(relax.op.divide(inp_1, inp_2)) + lhs, rhs = self._true_division_operands(lhs, rhs) + return self.block_builder.emit(relax.op.divide(lhs, rhs)) + + # With a rounding mode the result keeps the promoted dtype: an integer pair + # stays integer. Both operands are promoted first so a Python scalar lands in + # the tensor's dtype (or widens it) instead of arriving as a dtype-less + # constant that fails the same-dtype check on every float tensor. + lhs, rhs = self._promote_binary_operands(lhs, rhs) + if rounding_mode == "floor": + return self.block_builder.emit(relax.op.floor_divide(lhs, rhs)) + if rounding_mode == "trunc": + dtype = getattr(getattr(lhs, "ty", None), "dtype", None) + if dtype is not None and dtype.matches_code(DataTypeCode.INT, DataTypeCode.UINT): + # Integer division in relax truncates toward zero already. + return self.block_builder.emit(relax.op.divide(lhs, rhs)) + true_div = self.block_builder.emit(relax.op.divide(lhs, rhs)) return self.block_builder.emit(relax.op.trunc(true_div)) - else: - raise ValueError(f"Unsupported rounding_mode: {rounding_mode}") + raise ValueError(f"Unsupported rounding_mode: {rounding_mode}") def _fmod(self, node: fx.Node): args = self.retrieve_args(node) diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py b/python/tvm/relax/frontend/torch/exported_program_translator.py index 86c936723bf2..465185424bfd 100644 --- a/python/tvm/relax/frontend/torch/exported_program_translator.py +++ b/python/tvm/relax/frontend/torch/exported_program_translator.py @@ -98,10 +98,6 @@ def _log1p(self, node: fx.Node) -> relax.Var: one = relax.const(1, x.ty.dtype.dtype) return self.block_builder.emit(relax.op.log(relax.op.add(x, one))) - def _reciprocal(self, node: fx.Node) -> relax.Var: - x = self.env[node.args[0]] - return self.block_builder.emit(relax.op.divide(relax.const(1.0, x.ty.dtype.dtype), x)) - def _sqrt(self, node: fx.Node) -> relax.Var: x = self.env[node.args[0]] dtype = x.ty.dtype.dtype @@ -1850,8 +1846,8 @@ def create_convert_map( "bitwise_xor.Scalar": self._binary_op(relax.op.bitwise_xor, operator.xor), "bitwise_or_.Tensor": self._binary_op(relax.op.bitwise_or, operator.or_), "bitwise_or.Tensor": self._binary_op(relax.op.bitwise_or, operator.or_), - "div.Scalar": self._binary_op(relax.op.divide, operator.truediv), - "div.Tensor": self._binary_op(relax.op.divide, operator.truediv), + "div.Scalar": self._true_divide, + "div.Tensor": self._true_divide, "div.Tensor_mode": self._div, "eq.Scalar": self._binary_op(relax.op.equal, operator.eq), "eq.Tensor": self._binary_op(relax.op.equal, operator.eq), diff --git a/python/tvm/relax/frontend/torch/fx_translator.py b/python/tvm/relax/frontend/torch/fx_translator.py index 2e35ce6ce704..9df3438c1d76 100644 --- a/python/tvm/relax/frontend/torch/fx_translator.py +++ b/python/tvm/relax/frontend/torch/fx_translator.py @@ -64,10 +64,6 @@ def _fetch_attr(self, model, target: str): ########## Unary Ops ########## - def _reciprocal(self, node: fx.Node) -> relax.Var: - x = self.env[node.args[0]] - return self.block_builder.emit(relax.op.divide(relax.const(1.0, x.ty.dtype), x)) - def _leakyrelu_module(self, node: fx.Node) -> relax.Var: x = self.env[node.args[0]] module = self.named_modules[node.target] @@ -942,7 +938,7 @@ def create_convert_map( "rshift": self._binary_op(relax.op.right_shift, operator.rshift), "rsub": self._rsub, "sub": self._binary_op(relax.op.subtract, operator.sub), - "truediv": self._binary_op(relax.op.divide, operator.truediv), + "truediv": self._true_divide, "xor": self._binary_op(relax.op.bitwise_xor, operator.xor), # neural network "adaptive_avg_pool1d": self._adaptive_avg_pool1d, diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index 28ef877fd9cb..ad0b356525c7 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/tests/python/relax/test_frontend_from_exported_program.py @@ -1532,6 +1532,80 @@ def forward(self, x): _verify_scalar_promotion(model, x) +def test_true_division_of_integers_gives_float(): + # torch's `/` always produces a floating result: int64 / 2 is float32, not a + # truncating integer quotient. That rule sits on top of scalar promotion (an int + # scalar alone would not widen an int tensor), so it has its own converter. + class Div(Module): + def forward(self, x): + return x / 2 + + @tvm.script.ir_module + class expected_div: + @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((), dtype="float32") = R.astype(R.const(2, "int64"), dtype="float32") + lv2: R.Tensor((3,), dtype="float32") = R.divide(lv, lv1) + gv: R.Tuple(R.Tensor((3,), dtype="float32")) = (lv2,) + R.output(gv) + return gv + + verify_model(Div(), (torch.tensor([3, 4, 5]),), {}, expected_div) + + +@pytest.mark.parametrize( + "dtype, scalar", + [(torch.int64, 2), (torch.int32, 3), (torch.uint8, 2), (torch.bool, 2), (torch.int64, 2.5)], +) +def test_true_division_values(dtype, scalar): + class Div(Module): + def forward(self, x): + return x / scalar + + class RDiv(Module): + def forward(self, x): + return scalar / x + + class DivTensor(Module): + def forward(self, x): + return x / (x + 1) + + x = ( + torch.tensor([3, 4, 5], dtype=dtype) + if dtype is not torch.bool + else torch.tensor([True, True, False]) + ) + for model in (Div(), DivTensor()) + ((RDiv(),) if dtype is not torch.bool else ()): + _verify_scalar_promotion(model, x) + + +@pytest.mark.parametrize("dtype", [torch.int64, torch.int32, torch.float32]) +def test_division_with_rounding_mode(dtype): + # `//` and torch.div(..., rounding_mode=...) keep the promoted dtype -- an integer + # pair stays integer -- and the negative inputs tell floor and trunc apart. + class FloorDiv(Module): + def forward(self, x): + return x // 2 + + class DivFloor(Module): + def forward(self, x): + return torch.div(x, 2, rounding_mode="floor") + + class DivTrunc(Module): + def forward(self, x): + return torch.div(x, 2, rounding_mode="trunc") + + class DivFloorTensor(Module): + def forward(self, x): + return torch.div(x, x - 4, rounding_mode="floor") + + x = torch.tensor([-7, -3, 3, 7], dtype=dtype) + for model in (FloorDiv(), DivFloor(), DivTrunc(), DivFloorTensor()): + _verify_scalar_promotion(model, x) + + operator_binary_2 = [ (operator.eq, R.equal), (operator.ne, R.not_equal),