Repository: pytorch/pytorch

Issue:
Weird semantics for torch.diag()
The semantics of the function `torch.diag()` changes completely depending on the size of the inputs.
It the input is a vector, it creates a matrix with that vector as diagonal. It the input is a matrix, it extracts it's diagonal and returns a vector.
This is pretty weird to me, those should be two different functions. Also this doesn't consider a more than bi-dimensional case.

Ideally one function would take a tensor of any dimension and return its diagonal. This could be called `extract_diag`.
The other function should take a vector and a number of dimensions and should create a tensor with the number of dimensions specified with that vector as its diagonal. It should include also the actual offset parameter diagonal, maybe calling it offset would be easy to understand. This could be called `diag`.


Discussion:
For numpy users, this is exactly what they expect from [diag](https://docs.scipy.org/doc/numpy/reference/generated/numpy.diag.html).
Yeah, it's weird. I would not do it this way if I were designing an API from scratch. But a lot of our users come from NumPy or Matlab and they both have this behavior.

Maybe we should add `diagonal` and `diagflat`:
https://docs.scipy.org/doc/numpy/reference/generated/numpy.diagonal.html
https://docs.scipy.org/doc/numpy/reference/generated/numpy.diagflat.html 


I i think that that’s a weird choice in numpy too.
I understand trying to keep as much compibility to numpy as possible, but at the same time additional cleaner functions would be a nice idea. Call it diagonal or n_diagonal or other ways, but faithfullness to numpy shouldn’t be keep to have more reasonable functions imho.

Resolving patch:
diff --git a/torch/_torch_docs.py b/torch/_torch_docs.py
--- a/torch/_torch_docs.py
+++ b/torch/_torch_docs.py
@@ -1062,7 +1062,14 @@
     diagonal (int, optional): the diagonal to consider
     out (Tensor, optional): the output tensor
 
-Example:
+.. seealso::
+
+        :func:`torch.diagonal` always returns the diagonal of its input.
+
+        :func:`torch.diagflat` always constructs a tensor with diagonal elements
+        specified by the input.
+
+Examples:
 
 Get the square matrix where the input vector is the diagonal::
 
@@ -1115,6 +1122,109 @@
 
 """)
 
+add_docstr(torch.diagflat,
+           r"""
+diagflat(input, diagonal=0) -> Tensor
+
+- If :attr:`input` is a vector (1-D tensor), then returns a 2-D square tensor
+  with the elements of :attr:`input` as the diagonal.
+- If :attr:`input` is a tensor with more than one dimension, then returns a
+  2-D tensor with diagonal elements equal to a flattened :attr:`input`.
+
+The argument :attr:`offset` controls which diagonal to consider:
+
+- If :attr:`offset` = 0, it is the main diagonal.
+- If :attr:`offset` > 0, it is above the main diagonal.
+- If :attr:`offset` < 0, it is below the main diagonal.
+
+Args:
+    input (Tensor): the input tensor
+    offset (int, optional): the diagonal to consider. Default: 0 (main
+        diagonal).
+
+Examples:
+
+    >>> a = torch.randn(3)
+    >>> a
+
+     1.0480
+    -2.3405
+    -1.1138
+    [torch.FloatTensor of size 3]
+
+    >>> torch.diagflat(a)
+
+     1.0480  0.0000  0.0000
+     0.0000 -2.3405  0.0000
+     0.0000  0.0000 -1.1138
+    [torch.FloatTensor of size 3x3]
+
+    >>> torch.diagflat(a, 1)
+
+     0.0000  1.0480  0.0000  0.0000
+     0.0000  0.0000 -2.3405  0.0000
+     0.0000  0.0000  0.0000 -1.1138
+     0.0000  0.0000  0.0000  0.0000
+    [torch.FloatTensor of size 4x4]
+
+    >>> a = torch.randn(2, 2)
+    >>> a
+
+     0.1761 -0.9121
+    -0.5722  1.5219
+    [torch.FloatTensor of size (2,2)]
+
+    >>> torch.diagflat(a)
+
+     0.1761  0.0000  0.0000  0.0000
+     0.0000 -0.9121  0.0000  0.0000
+     0.0000  0.0000 -0.5722  0.0000
+     0.0000  0.0000  0.0000  1.5219
+    [torch.FloatTensor of size (4,4)]
+
+""")
+
+add_docstr(torch.diagonal,
+           r"""
+diagonal(input, offset=0) -> Tensor
+
+Returns a 1-D tensor with the diagonal elements of :attr:`input`.
+
+The argument :attr:`offset` controls which diagonal to consider:
+
+- If :attr:`offset` = 0, it is the main diagonal.
+- If :attr:`offset` > 0, it is above the main diagonal.
+- If :attr:`offset` < 0, it is below the main diagonal.
+
+Args:
+    input (Tensor): the input tensor. Must be 2-dimensional.
+    offset (int, optional): which diagonal to consider. Default: 0
+        (main diagonal).
+
+Examples:
+    >>> a = torch.randn(3, 3)
+    >>> a
+
+    -1.5328 -1.3210 -1.5204
+     0.8596  0.0471 -0.2239
+    -0.6617  0.0146 -1.0817
+    [torch.FloatTensor of size 3x3]
+
+    >>> torch.diagonal(a, 0)
+
+    -1.5328
+     0.0471
+    -1.0817
+    [torch.FloatTensor of size 3]
+
+    >>> torch.diagonal(a, 1)
+
+    -1.3210
+    -0.2239
+    [torch.FloatTensor of size 2]
+
+""")
+
 add_docstr(torch.dist,
            r"""
 dist(input, other, p=2) -> float


Test patch:
diff --git a/test/test_cuda.py b/test/test_cuda.py
--- a/test/test_cuda.py
+++ b/test/test_cuda.py
@@ -1517,6 +1517,12 @@ def test_arange(self):
             torch.arange(0, 10, out=b)
             self.assertEqual(a, b.cuda())
 
+    def test_diagonal(self):
+        TestTorch._test_diagonal(self, dtype=torch.cuda.float32)
+
+    def test_diagflat(self):
+        TestTorch._test_diagflat(self, dtype=torch.cuda.float32)
+
     @unittest.skipIf(torch.cuda.device_count() < 2, "only one GPU detected")
     def test_get_set_rng_state_all(self):
         states = torch.cuda.get_rng_state_all()
diff --git a/test/test_torch.py b/test/test_torch.py
--- a/test/test_torch.py
+++ b/test/test_torch.py
@@ -1198,6 +1198,51 @@ def test_diag(self):
         torch.diag(x, out=res2)
         self.assertEqual(res1, res2)
 
+    @staticmethod
+    def _test_diagonal(self, dtype=torch.float32):
+        x = torch.randn((100, 100), dtype=dtype)
+        result = torch.diagonal(x)
+        expected = torch.diag(x)
+        self.assertEqual(result, expected)
+
+        x = torch.randn((100, 100), dtype=dtype)
+        result = torch.diagonal(x, 17)
+        expected = torch.diag(x, 17)
+        self.assertEqual(result, expected)
+
+    def test_diagonal(self):
+        self._test_diagonal(self, dtype=torch.float32)
+
+    @staticmethod
+    def _test_diagflat(self, dtype=torch.float32):
+        # Basic sanity test
+        x = torch.randn((100,), dtype=dtype)
+        result = torch.diagflat(x)
+        expected = torch.diag(x)
+        self.assertEqual(result, expected)
+
+        # Test offset
+        x = torch.randn((100,), dtype=dtype)
+        result = torch.diagflat(x, 17)
+        expected = torch.diag(x, 17)
+        self.assertEqual(result, expected)
+
+        # Test where input has more than one dimension
+        x = torch.randn((2, 3, 4), dtype=dtype)
+        result = torch.diagflat(x)
+        expected = torch.diag(x.contiguous().view(-1))
+        self.assertEqual(result, expected)
+
+        # Noncontig input
+        x = torch.randn((2, 3, 4), dtype=dtype).transpose(2, 0)
+        self.assertFalse(x.is_contiguous())
+        result = torch.diagflat(x)
+        expected = torch.diag(x.contiguous().view(-1))
+        self.assertEqual(result, expected)
+
+    def test_diagflat(self):
+        self._test_diagflat(self, dtype=torch.float32)
+
     def test_eye(self):
         res1 = torch.eye(100, 100)
         res2 = torch.Tensor()