Enforce tensor a dtype == tensor b dtype for where.out in facto

Differential Revision: D82577515

Pull Request resolved: pytorch#14352

Author: ethansfng

backends/cadence/utils/facto_util.py

@@ -167,7 +167,25 @@ def random_size_constraint(deps: object, r: int, d: int) -> int:
                     cp.Size.Ge(lambda deps, r, d: 1),
                     max_size_constraint,
                 ]
-            else:
+            elif index == 1:  # input tensor(a)
+                tensor_constraints = [
+                    cp.Dtype.In(
+                        lambda deps: [
+                            torch.int8,
+                            torch.int16,
+                            torch.uint8,
+                            torch.uint16,
+                            torch.int32,
+                            torch.float32,
+                        ]
+                    ),
+                    cp.Value.Ge(lambda deps, dtype, struct: -(2**4)),
+                    cp.Value.Le(lambda deps, dtype, struct: 2**4),
+                    cp.Rank.Ge(lambda deps: 1),
+                    cp.Size.Ge(lambda deps, r, d: 1),
+                    max_size_constraint,
+                ]
+            else:  # input tensor(b)
                 tensor_constraints = [
                     cp.Dtype.In(
                         lambda deps: [
@@ -179,6 +197,7 @@ def random_size_constraint(deps: object, r: int, d: int) -> int:
                             torch.float32,
                         ]
                     ),
+                    cp.Dtype.Eq(lambda deps: deps[1].dtype),
                     cp.Value.Ge(lambda deps, dtype, struct: -(2**4)),
                     cp.Value.Le(lambda deps, dtype, struct: 2**4),
                     cp.Rank.Ge(lambda deps: 1),