migrate exir test to use torch.cond instead of the hop directly (#12695)

Summary:

torch.cond is the use-facing api, which allows closure, have safety checks etc.

Differential Revision: D78696006

Author: ydwu4

exir/backend/test/test_backends.py

@@ -1033,7 +1033,7 @@ def false_fn(x, y):
 
         def f(x, y):
             x = x + y
-            x = torch.ops.higher_order.cond(x[0][0] == 1, true_fn, false_fn, [x, y])
+            x = torch.cond(x[0][0] == 1, true_fn, false_fn, [x, y])
             x = x - y
             return x
 

exir/tests/control_flow_models.py

@@ -20,9 +20,7 @@ def true_branch(x):
         def false_branch(x):
             return x * x
 
-        return torch.ops.higher_order.cond(
-            inp.sum() > 4, true_branch, false_branch, [inp]
-        )
+        return torch.cond(inp.sum() > 4, true_branch, false_branch, [inp])
 
     def get_random_inputs(self):
         return (torch.rand(5),)
@@ -39,9 +37,7 @@ def true_branch(x):
         def false_branch(x):
             return x * x * x
 
-        return torch.ops.higher_order.cond(
-            inp.sum() > 4, true_branch, false_branch, [inp]
-        )
+        return torch.cond(inp.sum() > 4, true_branch, false_branch, [inp])
 
     def get_upper_bound_inputs(self):
         return (torch.rand(8),)
@@ -72,9 +68,7 @@ def true_branch(x):
         def false_branch(x):
             return x * 2
 
-        return torch.ops.higher_order.cond(
-            inp.sum() > 4, true_branch, false_branch, [inp]
-        )
+        return torch.cond(inp.sum() > 4, true_branch, false_branch, [inp])
 
     def get_random_inputs(self):
         return (torch.eye(5) * 2,)

exir/tests/test_passes.py

@@ -1463,9 +1463,7 @@ def forward(self, pred, x):
                 out = torch.nn.functional.linear(
                     x, self.w.to(torch.float16).to(torch.float32)
                 )
-                return torch.ops.higher_order.cond(
-                    pred, self.true_fn, self.false_fn, [out]
-                )
+                return torch.cond(pred, self.true_fn, self.false_fn, [out])
 
         mod = Module()
         x = torch.randn([3, 3])