improve tests and doc comment. Thank you, Adam!

t-vi · t-vi · commit a7d0b1fa87b6 · 2018-04-23T16:03:27.000+02:00
diff --git a/test/test_autograd.py b/test/test_autograd.py
@@ -2116,6 +2116,18 @@ def test_mul_out_result_requires_grad(self):
         # we should throw an exception if the output requires grad
         self.assertRaisesRegex(RuntimeError, 'out=', lambda: torch.mul(a, b, out=x))
 
+    def test_diagonal_derivative_requires_grad(self):
+        # test that the backward requires grad
+        # we do this is because diagonal_backward uses inplace
+        # operations and gradgradcheck does not catch whether
+        # they works as expected (it will succeed even if
+        # the gradient has requires_grad == False
+        a = torch.randn(5, 6, requires_grad=True)
+        b = torch.diagonal(a)**2
+        c = b.sum()
+        d, = torch.autograd.grad(c, a, retain_graph=True, create_graph=True)
+        self.assertTrue(d.requires_grad)
+
 
 def index_variable(shape, max_indices):
     if not isinstance(shape, tuple):
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -1913,36 +1913,20 @@ def test_diagonal(self):
     def test_diagonal_multidim(self):
         x = torch.randn(10, 11, 12, 13)
         xn = x.numpy()
-        result = torch.diagonal(x, 2, 2, 3)
-        expected = xn.diagonal(2, 2, 3)
-        self.assertEqual(expected.shape, result.shape)
-        self.assertTrue(np.allclose(expected, result.numpy()))
-        result = torch.diagonal(x, 2)
-        expected = torch.diagonal(x, 2, 0, 1)
-        self.assertEqual(expected, result)
-        result = torch.diagonal(x, -2, 1, 2)
-        expected = xn.diagonal(-2, 1, 2)
-        self.assertEqual(expected.shape, result.shape)
-        self.assertTrue(np.allclose(expected, result.numpy()))
-        result = torch.diagonal(x, 0, -2, -1)
-        expected = xn.diagonal(0, -2, -1)
-        self.assertEqual(expected.shape, result.shape)
-        self.assertTrue(np.allclose(expected, result.numpy()))
+        for args in [(2, 2, 3),
+                     (2,),
+                     (-2, 1, 2),
+                     (0, -2, -1)]:
+            result = torch.diagonal(x, *args)
+            expected = xn.diagonal(*args)
+            self.assertEqual(expected.shape, result.shape)
+            self.assertTrue(np.allclose(expected, result.numpy()))
         # test non-continguous
         xp = x.permute(1, 2, 3, 0)
         result = torch.diagonal(xp, 0, -2, -1)
         expected = xp.numpy().diagonal(0, -2, -1)
         self.assertEqual(expected.shape, result.shape)
         self.assertTrue(np.allclose(expected, result.numpy()))
-        # test that the backward requires grad
-        # we do this is because diagonal_backward uses inplace
-        # operations and gradgradcheck does not catch whether
-        # they works as expected
-        a = torch.randn(5, 6, requires_grad=True)
-        b = torch.diagonal(a)**2
-        c = b.sum()
-        d, = torch.autograd.grad(c,a, retain_graph=True, create_graph=True)
-        self.assertTrue(d.requires_grad)
 
     @staticmethod
     def _test_diagflat(self, dtype, device):
diff --git a/torch/_torch_docs.py b/torch/_torch_docs.py
@@ -1313,8 +1313,8 @@
     -0.2239
     [torch.FloatTensor of size 2]
 
-    >>> x = torch.randn(2,5,4,2)
-    >>> torch.diagonal(x, -1, 1, 2)
+    >>> x = torch.randn(2, 5, 4, 2)
+    >>> torch.diagonal(x, offset=-1, dim1=1, dim2=2)
 
     (0 ,.,.) =
      -0.6806 -0.0281 -0.6595 -0.4199
