Skip to content

[Fix][Relax][Frontend][Torch] Promote the tensor, not the scalar, in binary ops with a Python scalar #20372

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
hiyufan wants to merge 1 commit into
base: main
Choose a base branch
from
+159 −14
Open

[Fix][Relax][Frontend][Torch] Promote the tensor, not the scalar, in binary ops with a Python scalar #20372

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 48 additions & 4 deletions python/tvm/relax/frontend/torch/base_fx_graph_translator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 <op> 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."""
Expand Down Expand Up @@ -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))
Expand All @@ -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
Expand Down
121 changes: 111 additions & 10 deletions tests/python/relax/test_frontend_from_exported_program.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down Expand Up @@ -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

Expand Down
Loading