Repository: pytorch/pytorch

Issue:
[feature request] enhance diagonal functionality
# Objective
PyTorch currently does not have a way to take diagonals with respect to arbitrary axes as e.g. `numpy.diagonal` has.
Two immediate use cases are
- Taking the diagonal of a batch of matrices. This is handy e.g. in Gaussian processes, but I have seen this a number of times.
- One of the missing `einsum` features (`torch.einsum('jii->ji', a)`) would benefit from this as well.
# Plan
I propose to do the following (see https://github.com/t-vi/pytorch/tree/diagonal_with_dim for a preliminary implementation)
- Implement `diagonal` natively in ATen instead of referring to TH/THC `diag`.
- Make `diagonal` return a view by adjusting shape, stride, and offset. `diag` copies, but this seems impractical for higher dimensional arrays and also unneeded. *This is not backward compatible.*
- Add two dimension parameters dim1 and dim2.
- Make a diagonal_backward method.
- Include the Tensor method, too.

*Edit*: Based on your feedback, I updated the code to use numpy semantics.

# ~~Limiting numpy compatibility?~~
~~NumPy's implementation has two features that might be worth to differ from:~~
- ~~NumPy's default axes are 0 and 1. I think it is more natural to use -2 and -1 ("batch thinking").~~
- ~~NumPy moves the new axis replacing the two old ones at the end of the tensor. I think it might be more natural to put it into the place of the first axis to be removed. I must admit that this is more a gut feeling than based on hard facts.~~

I would greatly appreciate your input on this, in particular regarding the potential numpy deviations.


Discussion:

Resolving patch:
diff --git a/tools/autograd/gen_autograd.py b/tools/autograd/gen_autograd.py
--- a/tools/autograd/gen_autograd.py
+++ b/tools/autograd/gen_autograd.py
@@ -19,7 +19,7 @@
 deprecated_path = os.path.join(os.path.dirname(__file__), 'deprecated.yaml')
 
 VIEW_FUNCTIONS = {
-    'alias', 'as_strided', 'expand', 'narrow', 'permute', 'select', 'slice',
+    'alias', 'as_strided', 'diagonal', 'expand', 'narrow', 'permute', 'select', 'slice',
     'squeeze', 't', 'transpose', 'unfold', 'unsqueeze', 'view',
 }
 
diff --git a/torch/_torch_docs.py b/torch/_torch_docs.py
--- a/torch/_torch_docs.py
+++ b/torch/_torch_docs.py
@@ -1267,9 +1267,11 @@
 
 add_docstr(torch.diagonal,
            r"""
-diagonal(input, offset=0) -> Tensor
+diagonal(input, offset=0, dim1=0, dim2=1) -> Tensor
 
-Returns a 1-D tensor with the diagonal elements of :attr:`input`.
+Returns a partial view of :attr:`input` with the its diagonal elements
+with respect to :attr:`dim1` and :attr:`dim2` appended as a dimension
+at the end of the shape.
 
 The argument :attr:`offset` controls which diagonal to consider:
 
@@ -1278,9 +1280,15 @@
 - If :attr:`offset` < 0, it is below the main diagonal.
 
 Args:
-    input (Tensor): the input tensor. Must be 2-dimensional.
+    input (Tensor): the input tensor. Must be at least 2-dimensional.
     offset (int, optional): which diagonal to consider. Default: 0
         (main diagonal).
+    dim1 (int, optional): first dimension with respect to which to
+        take diagonal. Default: 0.
+    dim2 (int, optional): second dimension with respect to which to
+        take diagonal. Default: 1.
+
+.. note::  To take a batch diagonal, pass in dim1=-2, dim2=-1.
 
 Examples::
 
@@ -1305,6 +1313,17 @@
     -0.2239
     [torch.FloatTensor of size 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
+      0.8741 -0.1793 -0.6997  0.6265
+
+    (1 ,.,.) =
+      0.6182  1.3069  1.6503  1.7627
+     -0.2122 -0.2250  0.0990 -2.6433
+    [torch.FloatTensor of size (2,2,4)]
 """)
 
 add_docstr(torch.dist,


Test patch:
diff --git a/test/test_autograd.py b/test/test_autograd.py
--- 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):
@@ -2630,6 +2642,18 @@ class dont_convert(tuple):
     ('diag', (M,), NO_ARGS, '1d'),
     ('diag', (M, M), (1,), '2d_1'),
     ('diag', (M, M), (2,), '2d_2'),
+    ('diagonal', (M, M), NO_ARGS, '2d'),
+    ('diagonal', (3, 5), NO_ARGS, '2d_wide'),
+    ('diagonal', (3, 5), (2,), '2d_wide_pos'),
+    ('diagonal', (3, 5), (-2,), '2d_wide_neg'),
+    ('diagonal', (5, 3), NO_ARGS, '2d_tall'),
+    ('diagonal', (5, 3), (2,), '2d_tall_pos'),
+    ('diagonal', (5, 3), (-2,), '2d_tall_neg'),
+    ('diagonal', (M, M), (1,), '2d_1'),
+    ('diagonal', (M, M), (2,), '2d_2'),
+    ('diagonal', (M, M, M), (1, 1, 2), '3d_1'),
+    ('diagonal', (M, M, M), (2, 0, 1), '3d_2'),
+    ('diagonal', (M, M, M), (-2, 0, 1), '3d_3'),
     ('tril', (M, M), NO_ARGS),
     ('tril', (M, M), (2,), 'idx'),
     ('triu', (M, M), NO_ARGS),
diff --git a/test/test_torch.py b/test/test_torch.py
--- a/test/test_torch.py
+++ b/test/test_torch.py
@@ -1909,6 +1909,25 @@ def _test_diagonal(self, dtype, device):
     def test_diagonal(self):
         self._test_diagonal(self, dtype=torch.float32, device='cpu')
 
+    @unittest.skipIf(not TEST_NUMPY, 'Numpy not found')
+    def test_diagonal_multidim(self):
+        x = torch.randn(10, 11, 12, 13)
+        xn = x.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()))
+
     @staticmethod
     def _test_diagflat(self, dtype, device):
         # Basic sanity test