[Fix][Relax][Frontend][Torch] Make any and prod follow torch: missing overloads and result dtypes

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 ##########
 

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,

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