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

Summary:
Pull Request resolved: pytorch#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

Author: vkuzo

docs/source/torch.rst

@@ -536,3 +536,4 @@ Utilities
     set_deterministic
     is_deterministic
     vmap
+    Assert

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()

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()

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