From 9cb5ff9edac323ee9983e9fd328cd334daec8386 Mon Sep 17 00:00:00 2001 From: Chen Yufan Date: Thu, 17 Sep 2026 16:44:07 +0800 Subject: [PATCH] [Fix][Relax][Frontend][Torch] Make any and prod follow torch: missing overloads and result dtypes Two reductions were half-wired: - `any.default` (`x.any()`) had no converter. `any.dim` had one, but for anything other than a bool input it returned `max(x)` in the input dtype -- the largest value, not a truth value: `int32 [[0, 0, 5], [0, 0, 0]].any(1)` came back as int32 `[5, 0]` where torch gives bool `[True, False]`. torch.any asks whether any element is non-zero and is always bool. The converter now reduces the non-zero mask (`x != 0` for a non-bool input) with max in int8, since relax's max does not take bool, and casts back; `any.default` dispatches to it. - `prod.dim_int` (`x.prod(dim)`, with or without keepdim) had no converter, and `prod.default` kept the input dtype for a bool or integer input where torch accumulates in int64, exactly as `_sum` here already handles: `int32 [2^20, 2^20, 2] .prod()` overflowed to int32 where torch gives int64 2199023255552. `_prod` now applies the same rule as `_sum`, honours an explicit `dtype=`, treats an empty dim list as every axis, and serves both overloads. In the differential sweep of the frontend against torch.export (88 ops x 7 input shapes x static/dynamic), `any` goes from 14 raising programs to none and `prod` from 13 to one -- the 0-d input, which is a separate issue across all reductions. Tests: IR-level checks that `int32 x.prod(1)` casts to int64 before `R.prod` and that `int32 x.any(1)` lowers to not_equal -> int8 -> max -> bool; numeric checks of prod over every axis, one axis and keepdim (bool, int32 with an overflowing input, int64, float32), and of any in the same three forms over the same dtypes with an input whose second row is all zeros so both truth values appear. All six tests fail against the previous head. An explicit `dtype=` that already matches the input emits no cast, so `torch.prod(x, dtype=torch.float32)` on a float32 input still lowers to a bare `R.prod`, as `test_prod` in both frontend test files expects. --- .../torch/base_fx_graph_translator.py | 35 +++++-- .../torch/exported_program_translator.py | 2 + .../test_frontend_from_exported_program.py | 96 +++++++++++++++++++ 3 files changed, 125 insertions(+), 8 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..0ad97fc19f80 100644 --- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py +++ b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py @@ -1946,7 +1946,23 @@ def _prod(self, node: fx.Node) -> relax.Var: args = self.retrieve_args(node) x = args[0] dim = args[1] if len(node.args) > 1 else node.kwargs.get("dim", None) + if isinstance(dim, list | tuple) and len(dim) == 0: + dim = None keepdim = args[2] if len(node.args) > 2 else node.kwargs.get("keepdim", False) + dtype = node.kwargs.get("dtype", None) + if dtype is not None: + target = self._convert_data_type(dtype, self.env) + if str(x.ty.dtype) != str(target): + x = self.block_builder.emit(relax.op.astype(x, target)) + else: + # As for sum: torch accumulates a bool or integer product in int64 unless an + # explicit dtype is given. + input_dtype = x.ty.dtype.dtype + if input_dtype == "bool" or ( + (input_dtype.startswith("int") or input_dtype.startswith("uint")) + and input_dtype != "int64" + ): + x = self.block_builder.emit(relax.op.astype(x, "int64")) return self.block_builder.emit(relax.op.prod(x, dim, keepdims=keepdim)) def _std(self, node: fx.Node) -> relax.Var: @@ -2048,16 +2064,19 @@ def _any(self, node: fx.Node) -> relax.Var: args = self.retrieve_args(node) x = args[0] dim = args[1] if len(node.args) > 1 else node.kwargs.get("dim", None) + if isinstance(dim, list | tuple) and len(dim) == 0: + dim = None keepdim = args[2] if len(node.args) > 2 else node.kwargs.get("keepdim", False) - # max doesn't support boolean tensors directly, so we compute it in int8 and cast back - if x.ty.dtype == "bool": - x = relax.op.astype(x, "int8") - ret = relax.op.max(x, dim, keepdims=keepdim) - return self.block_builder.emit(relax.op.astype(ret, "bool")) - - # For boolean tensors, any is equivalent to max (checking if any element is True) - return self.block_builder.emit(relax.op.max(x, dim, keepdims=keepdim)) + # torch.any asks whether any element is non-zero and always returns bool. Reduce + # the non-zero mask with max in int8, since relax's max does not take bool, and + # cast back. Returning max(x) itself, as before, gave the input dtype and the + # largest value rather than a truth value for anything but a bool input. + if x.ty.dtype != "bool": + x = self.block_builder.emit(relax.op.not_equal(x, relax.const(0, x.ty.dtype))) + mask = self.block_builder.emit(relax.op.astype(x, "int8")) + ret = self.block_builder.emit(relax.op.max(mask, dim, keepdims=keepdim)) + return self.block_builder.emit(relax.op.astype(ret, "bool")) ########## Search ########## diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py b/python/tvm/relax/frontend/torch/exported_program_translator.py index 86c936723bf2..9de9eb628247 100644 --- a/python/tvm/relax/frontend/torch/exported_program_translator.py +++ b/python/tvm/relax/frontend/torch/exported_program_translator.py @@ -1960,11 +1960,13 @@ def create_convert_map( "upsample_nearest2d.vec": self._upsample_nearest2d, "upsample_bicubic2d.vec": self._upsample_bicubic2d, # statistical + "any.default": self._any, "any.dim": self._any, "any.dims": self._any, "mean.dim": self._mean, "mean.default": self._mean, "prod.default": self._prod, + "prod.dim_int": self._prod, "std.correction": self._std, "sum.default": self._sum, "sum.dim_IntList": self._sum, diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index d78629f5737b..9aec61422d69 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/tests/python/relax/test_frontend_from_exported_program.py @@ -7609,6 +7609,102 @@ def main( verify_model(Prod(), example_args, {}, Expected) +def test_prod_dim_and_integer_accumulation(): + # prod.dim_int had no converter, and prod on a bool or integer input kept the input + # dtype where torch accumulates in int64 (as sum already did here). + class ProdDim(Module): + def forward(self, x): + return torch.prod(x, 1) + + @tvm.script.ir_module + class expected_prod_dim: + @R.function + def main(x: R.Tensor((2, 3), dtype="int32")) -> R.Tuple(R.Tensor((2,), dtype="int64")): + with R.dataflow(): + lv: R.Tensor((2, 3), dtype="int64") = R.astype(x, dtype="int64") + lv1: R.Tensor((2,), dtype="int64") = R.prod(lv, axis=[1], keepdims=False) + gv: R.Tuple(R.Tensor((2,), dtype="int64")) = (lv1,) + R.output(gv) + return gv + + x = torch.tensor([[2**20, 2**20, 2], [1, 2, 3]], dtype=torch.int32) + verify_model(ProdDim(), (x,), {}, expected_prod_dim) + + +@pytest.mark.parametrize("dtype", [torch.bool, torch.int32, torch.int64, torch.float32]) +def test_prod_values(dtype): + class ProdAll(Module): + def forward(self, x): + return torch.prod(x) + + class ProdDim(Module): + def forward(self, x): + return torch.prod(x, 0) + + class ProdKeep(Module): + def forward(self, x): + return torch.prod(x, 1, keepdim=True) + + if dtype is torch.bool: + x = torch.tensor([[True, True, False], [True, True, True]]) + elif dtype is torch.int32: + x = torch.tensor([[2**20, 2**20, 2], [1, 2, 3]], dtype=dtype) # overflows int32 + else: + x = torch.tensor([[1, 2, 3], [4, 5, 6]]).to(dtype) + for model in (ProdAll(), ProdDim(), ProdKeep()): + with torch.no_grad(): + want = model(x) + mod = from_exported_program(export(model, (x,))) + assert str(mod["main"].ret_ty.fields[0].dtype) == str(want.dtype).replace("torch.", "") + verify_model_numerically(model, (x,)) + + +def test_any_returns_bool_for_every_dtype(): + # any.default had no converter, and any.dim on a non-bool input returned max(x) in + # the input dtype -- the largest value, not a truth value. torch.any is "is any + # element non-zero" and is always bool. + class AnyAll(Module): + def forward(self, x): + return torch.any(x) + + class AnyDim(Module): + def forward(self, x): + return torch.any(x, 1) + + class AnyKeep(Module): + def forward(self, x): + return torch.any(x, 0, keepdim=True) + + @tvm.script.ir_module + class expected_any_dim: + @R.function + def main(x: R.Tensor((2, 3), dtype="int32")) -> R.Tuple(R.Tensor((2,), dtype="bool")): + with R.dataflow(): + lv: R.Tensor((2, 3), dtype="bool") = R.not_equal(x, R.const(0, "int32")) + lv1: R.Tensor((2, 3), dtype="int8") = R.astype(lv, dtype="int8") + lv2: R.Tensor((2,), dtype="int8") = R.max(lv1, axis=[1], keepdims=False) + lv3: R.Tensor((2,), dtype="bool") = R.astype(lv2, dtype="bool") + gv: R.Tuple(R.Tensor((2,), dtype="bool")) = (lv3,) + R.output(gv) + return gv + + verify_model( + AnyDim(), (torch.tensor([[0, 0, 5], [0, 0, 0]], dtype=torch.int32),), {}, expected_any_dim + ) + + for dtype in (torch.bool, torch.int32, torch.int64, torch.float32): + if dtype is torch.bool: + x = torch.tensor([[False, False, True], [False, False, False]]) + else: + x = torch.tensor([[0, 0, 5], [0, 0, 0]]).to(dtype) + for model in (AnyAll(), AnyDim(), AnyKeep()): + with torch.no_grad(): + want = model(x) + mod = from_exported_program(export(model, (x,))) + assert str(mod["main"].ret_ty.fields[0].dtype) == "bool" + verify_model_numerically(model, (x,)) + + def test_cumprod(): class Cumprod(Module): def forward(self, x):