Fix bug with named tensors and (no) tracer support (#26106)

zou3519 · facebook-github-bot · commit babaac3e0853 · 2019-09-13T06:45:00.000-07:00
Summary: Pull Request resolved: #26106 Previously, in the named tensors build, an operator is marked as non-traceable if ANY of its overloads are named tensor overloads. This breaks the tracer for things like torch.full (has a names= overload for named tensor) and tensor.sum (has a Dimname overload for named tensor). This PR fixes the problem by putting the "no tracer support" logic into the location where the tracer attempts to construct a graph by adding a Dimname/DimnameList argument to a node. Test Plan: - new test in test_jit.py to check if torch.full is traceable - new test in test_namedtensor.py to check what happens when someone tries to trace a function that uses named tensor APIs. - [namedtensor ci] Differential Revision: D17353452 Pulled By: zou3519 fbshipit-source-id: b0b843c8357ffe54baee6e8df86db914f0b1ece4
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -1650,6 +1650,17 @@ def test_trace_arange(self):
     def test_trace_arange_with_grad(self):
         self.do_trace_arange(True)
 
+    # Test that a trace of torch.full(x.shape) doesn't store the shape as a constant
+    def test_trace_full_dynamic_shape(self):
+        def full_with_shape_like(x):
+            return torch.full(x.shape, 2)
+
+        x = torch.randn(3, 4)
+        ge = torch.jit.trace(full_with_shape_like, example_inputs=x)
+        y = torch.randn(2, 7)
+        self.assertEqual(ge(y).shape, y.shape)
+        self.assertEqual(ge(x).shape, x.shape)
+
     def test_trace_casts(self):
         casts = [
             lambda x: x.byte(),
diff --git a/test/test_namedtensor.py b/test/test_namedtensor.py
@@ -1088,7 +1088,22 @@ def test_as_strided(self):
     def test_as_strided_cuda(self):
         self._test_as_strided('cuda')
 
-    def test_no_jit_support(self):
+    def test_no_jit_tracer_support(self):
+        def foo(x):
+            return torch.full(x.shape, 2, names=('N',))
+
+        with self.assertRaisesRegex(RuntimeError, 'not supported with the tracer'):
+            x = torch.randn(3)
+            torch.jit.trace(foo, example_inputs=x)
+
+        def bar(x):
+            return x.select('N', 1)
+
+        with self.assertRaisesRegex(RuntimeError, 'not supported with the tracer'):
+            x = torch.randn(3)
+            torch.jit.trace(bar, example_inputs=x)
+
+    def test_no_jit_script_support(self):
         @torch.jit.script
         def foo(x):
             return x + 1
diff --git a/tools/autograd/gen_variable_type.py b/tools/autograd/gen_variable_type.py
@@ -281,11 +281,6 @@ def find_factory_functions(declarations):
 
 
 def should_trace(declaration):
-    # Short-term plan: Don't support tracing Dimname.
-    # Long-term plan: Add Dimname as a first-class type to the JIT.
-    if any('Dimname' in arg['simple_type'] for arg in declaration['arguments']):
-        return False
-
     # Operations involving Storage or Type are not traceable at the moment
     if any(arg['simple_type'] in {'Storage', 'Type', 'ConstQuantizerPtr'} for arg in declaration['arguments']):
         return False
diff --git a/torch/csrc/jit/tracer.cpp b/torch/csrc/jit/tracer.cpp
@@ -477,6 +477,14 @@ void addInputs(
     n->addInput(none);
   }
 }
+#ifdef BUILD_NAMEDTENSOR
+void addInputs(
+    Node* n,
+    const char* name,
+    c10::optional<at::DimnameList> value) {
+  TORCH_CHECK(false, "NYI: Named tensors are not supported with the tracer");
+}
+#endif
 void addInputs(
     Node* n,
     const char* name,
diff --git a/torch/csrc/jit/tracer.h b/torch/csrc/jit/tracer.h
@@ -4,6 +4,7 @@
 #include <c10/util/Exception.h>
 #include <torch/csrc/WindowsTorchApiMacro.h>
 #include <ATen/core/jit_type.h>
+#include <ATen/core/Dimname.h>
 
 #include <torch/csrc/utils/variadic.h>
 
@@ -280,6 +281,9 @@ TORCH_API void addInputs(
     const char* name,
     const c10::optional<at::ScalarType>& value);
 TORCH_API void addInputs(Node* n, const char* name, at::MemoryFormat value);
+#ifdef BUILD_NAMEDTENSOR
+TORCH_API void addInputs(Node* n, const char* name, c10::optional<at::DimnameList> value);
+#endif
 TORCH_API void addInputs(
     Node* n,
     const char* name,
