make the predicate truly data dependent. (#4130)

ydwu4 · facebook-github-bot · commit 5f2ab0e0d920 · 2024-07-08T15:22:44.000-07:00
Summary: Pull Request resolved: #4130 Make the pred truely data-dependent to avoid specializing into one of the branches if we land pytorch/pytorch#128709. Reviewed By: angelayi Differential Revision: D59287189 fbshipit-source-id: 0302cf385b441680cbacc3dafedef54f407ea0f1
diff --git a/exir/tests/control_flow_models.py b/exir/tests/control_flow_models.py
@@ -20,7 +20,7 @@ def false_branch(x):
             return x * x
 
         return torch.ops.higher_order.cond(
-            inp.size(0) > 4, true_branch, false_branch, [inp]
+            inp.sum() > 4, true_branch, false_branch, [inp]
         )
 
     def get_random_inputs(self):
@@ -39,7 +39,7 @@ def false_branch(x):
             return x * x * x
 
         return torch.ops.higher_order.cond(
-            inp.size(0) > 4, true_branch, false_branch, [inp]
+            inp.sum() > 4, true_branch, false_branch, [inp]
         )
 
     def get_upper_bound_inputs(self):
@@ -72,7 +72,7 @@ def false_branch(x):
             return x * 2
 
         return torch.ops.higher_order.cond(
-            inp.size(0) > 4, true_branch, false_branch, [inp]
+            inp.sum() > 4, true_branch, false_branch, [inp]
         )
 
     def get_random_inputs(self):

Original file line number	Diff line number	Diff line change
`@@ -20,7 +20,7 @@ def false_branch(x):`
`20`	`20`	`return x * x`
`21`	`21`
`22`	`22`	`return torch.ops.higher_order.cond(`
`23`		`- inp.size(0) > 4, true_branch, false_branch, [inp]`
	`23`	`+ inp.sum() > 4, true_branch, false_branch, [inp]`
`24`	`24`	`)`
`25`	`25`
`26`	`26`	`def get_random_inputs(self):`
`@@ -39,7 +39,7 @@ def false_branch(x):`
`39`	`39`	`return x * x * x`
`40`	`40`
`41`	`41`	`return torch.ops.higher_order.cond(`
`42`		`- inp.size(0) > 4, true_branch, false_branch, [inp]`
	`42`	`+ inp.sum() > 4, true_branch, false_branch, [inp]`
`43`	`43`	`)`
`44`	`44`
`45`	`45`	`def get_upper_bound_inputs(self):`
`@@ -72,7 +72,7 @@ def false_branch(x):`
`72`	`72`	`return x * 2`
`73`	`73`
`74`	`74`	`return torch.ops.higher_order.cond(`
`75`		`- inp.size(0) > 4, true_branch, false_branch, [inp]`
	`75`	`+ inp.sum() > 4, true_branch, false_branch, [inp]`
`76`	`76`	`)`
`77`	`77`
`78`	`78`	`def get_random_inputs(self):`