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):