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