Streamlining guard expect tests (#94404)
Changes:
* Add `simplified` kwarg to let you only render guards that are nontrivial (excludes duck sizing)
* Make a list of strings valid for sources, if you just have some variable names you want to bind to
* Add test helper `show_guards` using these facilities, switch a few tests to it
Signed-off-by: Edward Z. Yang <[email protected]>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/94404
Approved by: https://github.com/Chillee
diff --git a/test/test_proxy_tensor.py b/test/test_proxy_tensor.py
index 7368a85..b0170d4 100644
--- a/test/test_proxy_tensor.py
+++ b/test/test_proxy_tensor.py
@@ -12,7 +12,9 @@
from torch._subclasses.fake_tensor import DynamicOutputShapeException, DataDependentOutputException
from torch._decomp import decomposition_table
-from torch.fx.experimental.symbolic_shapes import sym_float, eval_guards, bind_symbols, fx_placeholder_vals
+from torch.fx.experimental.symbolic_shapes import (
+ sym_float, eval_guards, bind_symbols, fx_placeholder_vals, fx_placeholder_targets
+)
from torch.testing._internal.common_device_type import ops
from torch._C import _disabled_torch_function_impl
from torch.fx.experimental.proxy_tensor import make_fx, DecompositionInterpreter, get_isolated_graphmodule
@@ -35,6 +37,20 @@
HAS_CUDA = torch.cuda.is_available()
+def strip_end(s, suffix):
+ if suffix and s.endswith(suffix):
+ return s[:-len(suffix)]
+ else:
+ return s
+
+
+def show_guards(gm):
+ names = [strip_end(n, "_1") for n in fx_placeholder_targets(gm)]
+ return "\n".join(
+ gm.shape_env.produce_guards(fx_placeholder_vals(gm), names, simplified=True)
+ )
+
+
def process_failures():
"""
Takes file containing failures like
@@ -872,9 +888,7 @@
self.assertTrue(eval_guards(gm, torch.randn(4, 5)))
self.assertEqual(repr(bind_symbols(gm, torch.randn(4, 5))), "{s0: 4, s1: 5}")
self.assertFalse(eval_guards(gm, torch.randn(25, 5)))
- # TODO: There should eventually be guards for contiguity, but they're
- # not currently being done yet
- assert len(gm.shape_env.guards) == 1, "\n" + gm.shape_env.format_guards()
+ self.assertExpectedInline(show_guards(gm), """x.size()[0] < 20""")
@unittest.skipIf(not HAS_CUDA, 'CUDA-only test')
def test_cpu_scalar_cuda(self):
@@ -988,7 +1002,7 @@
gm = self._test_dynamic(f, [(1, 6), (8, 1)], test_inputs)
self.assertTrue(eval_guards(gm, torch.randn(1, 10), torch.randn(6, 1)))
self.assertFalse(eval_guards(gm, torch.randn(1, 2), torch.randn(4, 1)))
- assert len(gm.shape_env.guards) == 1
+ self.assertExpectedInline(show_guards(gm), """2*a.size()[1]*b.size()[0] > 20""")
def test_new_empty(self):
def f(a, b):
@@ -1156,7 +1170,7 @@
return final_vals
fx_g = _trace(f, 2, 4, 8, 16, 32)
- self._assert_no_guards(fx_g, 1)
+ self.assertExpectedInline(show_guards(fx_g), """""")
