From edf54723cbeaeddda7d8b8414a59f5d6137de8d4 Mon Sep 17 00:00:00 2001 From: Chen Yufan Date: Thu, 17 Sep 2026 16:24:43 +0800 Subject: [PATCH] [Fix][Relax][Frontend][Torch] Accumulate integer cumsum and cumprod in int64 With no `dtype` argument torch accumulates every integral and bool input of `cumsum` / `cumprod` in int64. The converters passed `dtype=None` through, so the running sum kept the input dtype and wrapped: uint8 [200, 100, 50].cumsum(1) torch int64 [200, 300, 350] before uint8 [200, 44, 94] int8 [100, 100, 50].cumsum(1) torch int64 [100, 200, 250] before int8 [100, -56, -6] int32 [2^30, 2^30, 5].cumsum(1) torch int64 [.., 2147483653] before int32 [.., -2147483643] bool [T, T, F].cumsum(0) torch int64 [1, 1, 0, ...] before InternalError An explicit `dtype=` was already honoured and is unchanged; float inputs keep their dtype, as in torch. The two converters share the rule through `_cumulative_dtype`. Swept cumsum / cumprod along both axes plus cumsum(dtype=float32) over bool, uint8, int8, int16, int32, int64, float16, float32 and float64 inputs chosen to overflow the narrow types, built with relax.build(llvm) and compared with torch on dtype and values: 25 matched / 16 wrong / 4 raised before, 45 / 0 / 0 after. Tests: an IR-level check that a uint8 cumsum emits `R.cumsum(..., dtype="int64")`, and numeric checks of cumsum, cumprod and cumsum(dtype=float32) over bool, uint8, int8, int32 and int64 inputs. The bool, uint8, int8 and int32 cases fail against the previous head; int64 passes there and pins that it is untouched. --- .../torch/base_fx_graph_translator.py | 30 ++++++----- .../test_frontend_from_exported_program.py | 53 +++++++++++++++++++ 2 files changed, 69 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..6f5ffdf4aba7 100644 --- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py +++ b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py @@ -2120,29 +2120,31 @@ def _chunk(self, node: fx.Node) -> relax.Var: relax.op.split(x=x, indices_or_sections=n_sections, axis=dim) ) - def _cumprod(self, node: fx.Node) -> relax.Var: - x = self.env[node.args[0]] + def _cumulative_dtype(self, x: relax.Expr, node: fx.Node) -> str | None: + """The output dtype of torch.cumsum / torch.cumprod for ``x``. - dim = node.args[1] if len(node.args) > 1 else node.kwargs.get("dim", None) + With no ``dtype`` argument torch accumulates every integral and bool input in + int64 -- ``uint8 [200, 100].cumsum(0)`` is int64 ``[200, 300]``, not a uint8 that + wraps to ``[200, 44]`` -- and keeps floating inputs as they are. + """ if "dtype" in node.kwargs: - dtype = self._convert_data_type(str(node.kwargs["dtype"]), self.env) - else: - dtype = None + return self._convert_data_type(str(node.kwargs["dtype"]), self.env) + dtype = x.ty.dtype + if dtype.matches_code(DataTypeCode.INT, DataTypeCode.UINT) or str(dtype) == "bool": + return "int64" + return None - return self.block_builder.emit(relax.op.cumprod(x, dim, dtype)) + def _cumprod(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", None) + return self.block_builder.emit(relax.op.cumprod(x, dim, self._cumulative_dtype(x, node))) def _cumsum(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", None) - if "dtype" in node.kwargs: - dtype = self._convert_data_type(str(node.kwargs["dtype"]), self.env) - else: - dtype = None if "out" in node.kwargs: raise ValueError("specifying out for cumsum is not supported yet") - - return self.block_builder.emit(relax.op.cumsum(x, dim, dtype)) + return self.block_builder.emit(relax.op.cumsum(x, dim, self._cumulative_dtype(x, node))) def _expand(self, node: fx.Node) -> relax.Var: args = self.retrieve_args(node) diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index d78629f5737b..fcfc41cca50f 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/tests/python/relax/test_frontend_from_exported_program.py @@ -5037,6 +5037,59 @@ def main(input_1: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tuple( verify_model(Cumsum(), example_args, {}, expected1) +def test_cumsum_integer_input_accumulates_in_int64(): + # With no dtype argument torch accumulates every integral and bool input in int64. + # Keeping the input dtype instead makes the running sum wrap: uint8 [200, 100, 50] + # is [200, 44, 94] rather than [200, 300, 350]. + class Cumsum(Module): + def forward(self, x): + return torch.cumsum(x, dim=1) + + @tvm.script.ir_module + class expected: + @R.function + def main(x: R.Tensor((2, 3), dtype="uint8")) -> R.Tuple(R.Tensor((2, 3), dtype="int64")): + with R.dataflow(): + lv: R.Tensor((2, 3), dtype="int64") = R.cumsum(x, axis=1, dtype="int64") + gv: R.Tuple(R.Tensor((2, 3), dtype="int64")) = (lv,) + R.output(gv) + return gv + + x = torch.tensor([[200, 100, 50], [2, 3, 4]], dtype=torch.uint8) + verify_model(Cumsum(), (x,), {}, expected) + + +@pytest.mark.parametrize("dtype", [torch.bool, torch.uint8, torch.int8, torch.int32, torch.int64]) +def test_cumsum_cumprod_integer_values(dtype): + class Cumsum(Module): + def forward(self, x): + return torch.cumsum(x, dim=1) + + class Cumprod(Module): + def forward(self, x): + return torch.cumprod(x, dim=1) + + class CumsumFloat(Module): + def forward(self, x): + return torch.cumsum(x, dim=0, dtype=torch.float32) + + if dtype is torch.bool: + x = torch.tensor([[True, True, False], [True, True, True]]) + elif dtype is torch.uint8: + x = torch.tensor([[200, 100, 50], [2, 3, 4]], dtype=dtype) + elif dtype is torch.int8: + x = torch.tensor([[100, 100, 50], [2, 3, 4]], dtype=dtype) + else: + x = torch.tensor([[2**30, 2**30, 5], [2, 3, 4]], dtype=dtype) + for model in (Cumsum(), Cumprod(), CumsumFloat()): + 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.", "") + verify_model_numerically(model, (x,)) + + def test_expand(): class Expand1(Module): def forward(self, x):