Some test cases for no_grad

Author: zou3519

test/test_eager_transforms.py

@@ -485,6 +485,119 @@ def foo(x):
                 tensor(3.1400)))""")
         self.assertEqual(buf, expected)
 
+    @unittest.expectedFailure
+    def test_no_grad_outside(self, device):
+        x = torch.randn([], device=device)
+        with torch.no_grad():
+            y = grad(torch.sin)(x)
+        self.assertEqual(y, x.cos())
+
+    def test_no_grad_inside(self, device):
+        def f(x):
+            with torch.no_grad():
+                shift = x ** 2
+            return x ** 2 - shift
+
+        x = torch.randn([], device=device)
+        y = grad(f)(x)
+        self.assertEqual(y, 2 * x)
+        y = grad(grad(f))(x)
+        self.assertEqual(y, 2)
+
+    @unittest.expectedFailure
+    def test_no_grad_mixed(self, device):
+        def f(x):
+            with torch.no_grad():
+                shift = x ** 2
+            return x ** 2 - shift
+
+        x = torch.randn([], device=device)
+        with torch.no_grad():
+            y = grad(f)(x)
+
+        self.assertEqual(y, 2 * x)
+
+    def test_no_grad_nested_simple(self, device):
+        def h(x):
+            with torch.no_grad():
+                shift = grad(lambda x: x ** 3)(x)
+            return x ** 3 - shift
+
+        x = torch.tensor(2., device=device)
+        y = grad(h)(x)
+        self.assertEqual(y, 6 * x)
+
+    def test_no_grad_nested_complicated(self, device):
+        def f(x):
+            with torch.no_grad():
+                shift = x ** 3
+            return x ** 3 - shift
+
+        def g(x):
+            r1 = grad(f)(x)
+            with torch.no_grad():
+                shift = grad(f)(x)
+            return r1 - shift
+
+        x = torch.randn([])
+        y = grad(g)(x)
+        # The only differential part of g is x ** 3
+        self.assertEqual(y, 6 * x)
+
+    def test_no_grad_value(self, device):
+        def h(x):
+            with torch.no_grad():
+                gvalue, value = grad_and_value(lambda x: x ** 3)(x)
+            return x ** 3 - value
+
+        x = torch.tensor(2., device=device)
+        y = grad(h)(x)
+        self.assertEqual(y, 6 * x)
+
+    @unittest.expectedFailure
+    def test_no_grad_outside_vjp(self, device):
+        def h(x):
+            return x ** 2
+
+        x = torch.tensor(2., requires_grad=True, device=device)
+        with torch.no_grad():
+            out, vjp_fn = vjp(h, x)
+            y, = vjp_fn(1.)
+
+        self.assertEqual(y, 2 * x)
+        self.assertFalse(y.requires_grad)
+        self.assertFalse(out.requires_grad)
+
+    @unittest.expectedFailure
+    def test_no_grad_outside_vjp_fn(self, device):
+        def h(x):
+            return x ** 2
+
+        x = torch.tensor(2., requires_grad=True, device=device)
+        out, vjp_fn = vjp(h, x)
+        with torch.no_grad():
+            y, = vjp_fn(1.)
+
+        self.assertEqual(y, 2 * x)
+        self.assertFalse(y.requires_grad)
+        self.assertTrue(out.requires_grad)
+
+    @unittest.expectedFailure
+    def test_no_grad_outside_vjp_only(self, device):
+        def h(x):
+            return x ** 2
+
+        x = torch.tensor(2., requires_grad=True, device=device)
+        with torch.no_grad():
+            out, vjp_fn = vjp(h, x)
+        y, = vjp_fn(1.)
+
+        self.assertEqual(y, 2 * x)
+        self.assertFalse(out.requires_grad)
+
+        # This one is a little weird. `vjp_fn` didn't save enough info
+        # during the forward pass for the output to be differentiable.
+        self.assertFalse(y.requires_grad)
 
 class TestVmapOfGrad(TestCase):
     def test_per_sample_grads_inplace_view(self, device):