support torch._C._get_tracing_state in script

Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/23248 Test Plan: Imported from OSS Differential Revision: D16466588 Pulled By: wanchaol fbshipit-source-id: 3c3d5dec2cea2f9cb080eadaef457cc62ac3fbe0

support torch._C._get_tracing_state in script
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/23248 Test Plan: Imported from OSS Differential Revision: D16466588 Pulled By: wanchaol fbshipit-source-id: 3c3d5dec2cea2f9cb080eadaef457cc62ac3fbe0
c384fbf4 · Wanchao Liang · Facebook Github Bot · e1f89859 · c384fbf4 · c384fbf4
--- a/test/test_jit.py
+++ b/test/test_jit.py
@@ -7809,6 +7809,17 @@ a")
        f = io.BytesIO()
        torch.onnx._export(m, (x, seq_lens), f, verbose=False)

+    def test_script_get_tracing_state(self):
+        def test_if_tracing(x):
+            if torch._C._get_tracing_state():
+                return x + 1
+            else:
+                return x - 1
+
+        inp = torch.randn(3, 3)
+
+        self.checkScript(test_if_tracing, (inp,))
+
    def test_script_outputs(self):
        with self.assertRaisesRegex(RuntimeError, "cannot be used as a tuple"):
            @torch.jit.script

--- a/torch/csrc/jit/register_special_ops.cpp
+++ b/torch/csrc/jit/register_special_ops.cpp
@@ -464,6 +464,13 @@ RegisterOperators reg({
        "Tensor? unsorted_indices) -> (Tensor, Tensor, Tensor?, Tensor?)",
        [](Stack& stack) { return 0; },
        aliasAnalysisFromSchema()),
+    Operator(
+        "aten::_get_tracing_state() -> bool",
+        [](Stack& stack) {
+          push(stack, false);
+          return 0;
+        },
+        aliasAnalysisFromSchema()),
    Operator(
        "aten::_no_grad_uniform_(Tensor(a!) tensor, float a, float b) -> Tensor(a!)",
        [](Stack& stack) {

--- a/torch/jit/__init__.py
+++ b/torch/jit/__init__.py
@@ -2003,6 +2003,7 @@ def _get_builtin_table():
        (torch.nn.init._no_grad_uniform_, "aten::_no_grad_uniform_"),
        (torch.nn.init._no_grad_zero_, "aten::_no_grad_zero_"),
        (torch.nn.utils.rnn.get_packed_sequence, "aten::_pack_sequence"),
+        (torch._C._get_tracing_state, "aten::_get_tracing_state"),
        (warnings.warn, "aten::warn"),
    ]