diag() and diagflat() in PyTorch
Tags: pytorch, diag, diagflat, diagonal
Published: 2024-05-22 19:56:40

*Memos:
- [My post](https://dev.to/hyperkai/eye-in-pytorch-1mjj) explains [eye()](https://pytorch.org/docs/stable/generated/torch.eye.html).
- [My post](https://dev.to/hyperkai/diagonal-in-pytorch-2643) explains [diagonal()](https://pytorch.org/docs/stable/generated/torch.diagonal.html).
- [My post](https://dev.to/hyperkai/diagembed-in-pytorch-a2o) explains [diag_embed()](https://pytorch.org/docs/stable/generated/torch.diag_embed.html).

[diag()](https://pytorch.org/docs/stable/generated/torch.diag.html) can create the 2D tensor of zero or more elements on the diagonal and zero or more `0`, `0.`, `0.+0.j` or `False` elsewhere from the 1D tensor of zero or more elements or can extract the 1D tensor of zero or more elements on the diagonal from the 2D tensor of zero or more elements as shown below:

*Memos:
- `diag()` can be used with [torch](https://pytorch.org/docs/stable/torch.html) or a tensor.
- The 1st argument with `torch` or using a tensor is `input`(Required-Type:`tensor` of `int`, `float`, `complex` or `bool`). *Only a 2D or 1D tensor can be used.
- The 2nd argument with `torch` or the 1st argument with a tensor is `diagonal`(Optional-Default:`0`-Type:`int`).
- There is `out` argument with torch(Optional-Type:`tensor`). *[My post](https://dev.to/hyperkai/set-out-with-out-argument-functions-pytorch-3ee) explains `out` argument.
- A 2D tensor creates a 1D tensor.
- A 1D tensor creates a 2D tensor.

```python
import torch

my_tensor = torch.tensor([7, -4, 5])

torch.diag(input=my_tensor)
my_tensor.diag()
torch.diag(input=my_tensor, diagonal=0)
# tensor([[7, 0, 0],
#         [0, -4, 0],
#         [0, 0, 5]])

torch.diag(input=my_tensor, diagonal=1)
# tensor([[0, 7, 0, 0],
#         [0, 0, -4, 0],
#         [0, 0, 0, 5],
#         [0, 0, 0, 0]])

torch.diag(input=my_tensor, diagonal=-1)
# tensor([[0, 0, 0, 0],
#         [7, 0, 0, 0],
#         [0, -4, 0, 0],
#         [0, 0, 5, 0]])

torch.diag(input=my_tensor, diagonal=2)
# tensor([[0, 0, 7, 0, 0],
#         [0, 0, 0, -4, 0],
#         [0, 0, 0, 0, 5],
#         [0, 0, 0, 0, 0],
#         [0, 0, 0, 0, 0]])

torch.diag(input=my_tensor, diagonal=-2)
# tensor([[0, 0, 0, 0, 0],
#         [0, 0, 0, 0, 0],
#         [7, 0, 0, 0, 0],
#         [0, -4, 0, 0, 0],
#         [0, 0, 5, 0, 0]])

my_tensor = torch.tensor([7., -4., 5.])

torch.diag(input=my_tensor)
# tensor([[7., 0., 0.],
#         [0., -4., 0.],
#         [0., 0., 5.]])

my_tensor = torch.tensor([7.+0.j, -4.+0.j, 5.+0.j])

torch.diag(input=my_tensor)
# tensor([[7.+0.j, 0.+0.j, 0.+0.j],
#         [0.+0.j, -4.+0.j, 0.+0.j],
#         [0.+0.j, 0.+0.j, 5.+0.j]])

my_tensor = torch.tensor([True, True, True])

torch.diag(input=my_tensor)
# tensor([[True, False, False],
#         [False, True, False],
#         [False, False, True]])

my_tensor = torch.tensor([[7, -4, 5],
                          [-6, -3, 8],
                          [9, 1, -2]])
torch.diag(input=my_tensor)
torch.diag(input=my_tensor, diagonal=0)
# tensor([7, -3, -2])

torch.diag(input=my_tensor, diagonal=1)
# tensor([-4, 8])

torch.diag(input=my_tensor, diagonal=-1)
# tensor([-6, -1])

torch.diag(input=my_tensor, diagonal=2)
# tensor([5])

torch.diag(input=my_tensor, diagonal=-2)
tensor([9])
```

[diagflat()](https://pytorch.org/docs/stable/generated/torch.diagflat.html) can create the 2D tensor of zero or more elements on the diagonal and zero or more `0`, `0.`, `0.+0.j` or `False` elsewhere from the 0D or more D tensor of zero or more elements as shown below:

*Memos:
- `diagflat()` can be used with `torch` or a tensor.
- The 1st argument with `torch` or using a tensor is `input`(Required-Type:`tensor` of `int`, `float`, `complex` or `bool`).
- The 2nd argument with `torch` or the 1st argument with a tensor is `offset`(Optional-Default:`0`-Type:`int`).

```python
import torch

my_tensor = torch.tensor([7, -4, 5])

torch.diagflat(input=my_tensor)
my_tensor.diagflat()
torch.diagflat(input=my_tensor, offset=0)
# tensor([[7, 0, 0],
#         [0, -4, 0],
#         [0, 0, 5]])

torch.diagflat(input=my_tensor, offset=1)
# tensor([[0, 7, 0, 0],
#         [0, 0, -4, 0],
#         [0, 0, 0, 5],
#         [0, 0, 0, 0]])

torch.diagflat(input=my_tensor, offset=-1)
# tensor([[0, 0, 0, 0],
#         [7, 0, 0, 0],
#         [0, -4, 0, 0],
#         [0, 0, 5, 0]])

torch.diagflat(input=my_tensor, offset=2)
# tensor([[0, 0, 7, 0, 0],
#         [0, 0, 0, -4, 0],
#         [0, 0, 0, 0, 5],
#         [0, 0, 0, 0, 0],
#         [0, 0, 0, 0, 0]])

torch.diagflat(input=my_tensor, offset=-2)
# tensor([[0, 0, 0, 0, 0],
#         [0, 0, 0, 0, 0],
#         [7, 0, 0, 0, 0],
#         [0, -4, 0, 0, 0],
#         [0, 0, 5, 0, 0]])

my_tensor = torch.tensor([7., -4., 5.])

torch.diagflat(input=my_tensor)
# tensor([[7., 0., 0.],
#         [0., -4., 0.],
#         [0., 0., 5.]])

my_tensor = torch.tensor([7.+0.j, -4.+0.j, 5.+0.j])

torch.diagflat(input=my_tensor)
# tensor([[7.+0.j, 0.+0.j, 0.+0.j],
#         [0.+0.j, -4.+0.j, 0.+0.j],
#         [0.+0.j, 0.+0.j, 5.+0.j]])

my_tensor = torch.tensor([True, True, True])

torch.diagflat(input=my_tensor)
# tensor([[True, False, False],
#         [False, True, False],
#         [False, False, True]])

my_tensor = torch.tensor([[7, -4, 5],
                          [-6, -3, 8],
                          [9, 1, -2]])
torch.diagflat(input=my_tensor)
torch.diagflat(input=my_tensor, offset=0)
# tensor([[7, 0, 0, 0, 0, 0, 0, 0, 0],
#         [0, -4, 0, 0, 0, 0, 0, 0, 0],
#         [0, 0, 5, 0, 0, 0, 0, 0, 0],
#         [0, 0, 0, -6, 0, 0, 0, 0, 0],
#         [0, 0, 0, 0, -3, 0, 0, 0, 0],
#         [0, 0, 0, 0, 0, 8, 0, 0, 0],
#         [0, 0, 0, 0, 0, 0, 9, 0, 0],
#         [0, 0, 0, 0, 0, 0, 0, 1, 0],
#         [0, 0, 0, 0, 0, 0, 0, 0, -2]])

torch.diagflat(input=my_tensor, offset=1)
# tensor([[0, 7, 0, 0, 0, 0, 0, 0, 0, 0],
#         [0, 0, -4, 0, 0, 0, 0, 0, 0, 0],
#         [0, 0, 0, 5, 0, 0, 0, 0, 0, 0],
#         [0, 0, 0, 0, -6, 0, 0, 0, 0, 0],
#         [0, 0, 0, 0, 0, -3, 0, 0, 0, 0],
#         [0, 0, 0, 0, 0, 0, 8, 0, 0, 0],
#         [0, 0, 0, 0, 0, 0, 0, 9, 0, 0],
#         [0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
#         [0, 0, 0, 0, 0, 0, 0, 0, 0, -2],
#         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])

torch.diagflat(input=my_tensor, offset=-1)
# tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
#         [7, 0, 0, 0, 0, 0, 0, 0, 0, 0],
#         [0,-4, 0, 0, 0, 0, 0, 0, 0, 0],
#         [0, 0, 5, 0, 0, 0, 0, 0, 0, 0],
#         [0, 0, 0, -6, 0, 0, 0, 0, 0, 0],
#         [0, 0, 0, 0, -3, 0, 0, 0, 0, 0],
#         [0, 0, 0, 0, 0, 8, 0, 0, 0, 0],
#         [0, 0, 0, 0, 0, 0, 9, 0, 0, 0],
#         [0, 0, 0, 0, 0, 0, 0, 1, 0, 0],
#         [0, 0, 0, 0, 0, 0, 0, 0, -2, 0]])

torch.diagflat(input=my_tensor, offset=2)
# tensor([[0, 0, 7, 0, 0, 0, 0, 0, 0, 0, 0],
#         [0, 0, 0, -4, 0, 0, 0, 0, 0, 0, 0],
#         [0, 0, 0, 0, 5, 0, 0, 0, 0, 0, 0],
#         [0, 0, 0, 0, 0, -6, 0, 0, 0, 0, 0],
#         [0, 0, 0, 0, 0, 0, -3, 0, 0, 0, 0],
#         [0, 0, 0, 0, 0, 0, 0, 8, 0, 0, 0],
#         [0, 0, 0, 0, 0, 0, 0, 0, 9, 0, 0],
#         [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
#         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -2],
#         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
#         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])

torch.diagflat(input=my_tensor, offset=-2)
# tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
#         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
#         [7, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
#         [0, -4, 0, 0, 0, 0, 0, 0, 0, 0, 0],
#         [0, 0, 5, 0, 0, 0, 0, 0, 0, 0, 0],
#         [0, 0, 0, -6, 0, 0, 0, 0, 0, 0, 0],
#         [0, 0, 0, 0, -3, 0, 0, 0, 0, 0, 0],
#         [0, 0, 0, 0, 0, 8, 0, 0, 0, 0, 0],
#         [0, 0, 0, 0, 0, 0, 9, 0, 0, 0, 0],
#         [0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
#         [0, 0, 0, 0, 0, 0, 0, 0, -2, 0, 0]])
```