Add torch.do_assert, which is symbolically traceable (#45188)

vkuzo · facebook-github-bot · commit eee7dad3768f · 2020-09-25T13:46:28.000-07:00
Summary: Pull Request resolved: #45188 This is a symbolically traceable alternative to Python's `assert`. It should be useful to allow people who want to use FX to also be able to assert things. A bunch of TODO(before) land are inline - would love thoughts on where is the best place for this code to live, and what this function should be called (since `assert` is reserved). Test Plan: ``` python test/test_fx.py TestFX.test_symbolic_trace_assert ``` Imported from OSS Reviewed By: jamesr66a Differential Revision: D23861567 fbshipit-source-id: d9d6b9556140faccc0290eba1fabea401d7850de
diff --git a/docs/source/torch.rst b/docs/source/torch.rst
@@ -536,3 +536,4 @@ Utilities
     set_deterministic
     is_deterministic
     vmap
+    Assert
diff --git a/test/test_fx.py b/test/test_fx.py
@@ -603,6 +603,22 @@ def test_construct_root_dict(self):
         ref_out : torch.Tensor = linear_mod(x) + add_param
         self.assertEqual(out, ref_out)
 
+    def test_symbolic_trace_assert(self):
+        message = "assert_foobar"
+
+        class AssertsTensorShape(torch.nn.Module):
+            def forward(self, x):
+                torch.Assert(x.shape[1] > 4, message)
+                return x
+
+        m = AssertsTensorShape()
+        # verify traceability
+        traced = symbolic_trace(m)
+        # verify assertion on traced model works correctly at runtime
+        traced(torch.rand(4, 5))
+        with self.assertRaisesRegex(AssertionError, message):
+            traced(torch.rand(4, 3))
+
 
 if __name__ == '__main__':
     run_tests()
diff --git a/test/test_utils.py b/test/test_utils.py
@@ -795,5 +795,13 @@ def test_fuzzer(self):
                 x, torch.Tensor(expected_results[i]), rtol=1e-3, atol=1e-3)
 
 
+class TestAssert(TestCase):
+    def test_assert_true(self):
+        # verify assertions work as expected
+        torch.Assert(True, "foo")
+        with self.assertRaisesRegex(AssertionError, "bar"):
+            torch.Assert(False, "bar")
+
+
 if __name__ == '__main__':
     run_tests()
diff --git a/torch/__init__.py b/torch/__init__.py
@@ -612,3 +612,12 @@ def compiled_with_cxx11_abi():
 # class usage. We add these lines here to preserve backward compatbility.
 quantized_lstm = torch.ops.aten.quantized_lstm
 quantized_gru = torch.ops.aten.quantized_gru
+
+from .overrides import has_torch_function, handle_torch_function
+
+def Assert(condition, message):
+    r"""A wrapper around Python's assert which is symbolically traceable.
+    """
+    if type(condition) is not torch.Tensor and has_torch_function((condition,)):
+        return handle_torch_function(Assert, (condition,), condition, message)
+    assert condition, message
