DEV Community

Super Kai (Kazuya Ito)
Super Kai (Kazuya Ito)

Posted on

diag() and diagflat() in PyTorch

*Memos:

diag() can create a 2D tensor with a 1D tensor on the diagonal and zero or more 0, 0., 0.+0.j or False elsewhere or extract a 1D tensor from a 2D tensor on the diagonal as shown below:

*Memos:

  • diag() can be used with torch or a tensor.
  • The 1st argument(tensor of int, float, complex or bool) with torch or using a tensor(tensor of int, float, complex or bool) is input(Required). *Only a 2D or 1D tensor can be used.
  • The 2nd argument(int) with torch or the 1st argument(int) with a tensor is diagonal(Optional-Default:0).
  • A 2D tensor creates a 1D tensor.
  • A 1D tensor creates a 2D tensor.
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])
Enter fullscreen mode Exit fullscreen mode

diagflat() can create a 2D tensor with a 0D or more D tensor on the diagonal and zero or more 0, 0., 0.+0.j or False elsewhere as shown below:

*Memos:

  • diagflat() can be used with torch or a tensor.
  • The 1st argument(tensor of int, float, complex or bool) with torch or using a tensor(tensor of int, float, complex or bool) is input(Required).
  • The 2nd argument(int) with torch or the 1st argument(int) with a tensor is offset(Optional-Default:0).
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]])
Enter fullscreen mode Exit fullscreen mode

Top comments (0)