diag_embed() in PyTorch
Tags: pytorch, diagembed, diagonal, function
Published: 2024-05-22 00:30:47

*Memos:
- [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/eye-in-pytorch-1mjj) explains [eye()](https://pytorch.org/docs/stable/generated/torch.eye.html).
- [My post](https://dev.to/hyperkai/diag-and-diagflat-in-pytorch-22n9) explains [diag()](https://pytorch.org/docs/stable/generated/torch.diag.html) and [diagflat()](https://pytorch.org/docs/stable/generated/torch.diagflat.html).

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

*Memos:
- `diag_embed()` 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`).
- The 2nd argument with torch or the 1st argument with a tensor is `offset`(Optional-Default:`0`-Type:`int`).
- The 3rd argument with torch or the 2nd argument with a tensor is `dim1`(Optional-Default:`-2`-Type:`int`).
- The 4th argument with torch or the 3rd argument with a tensor is `dim2`(Optional-Default:`-1`-Type:`int`).

```python
import torch

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

torch.diag_embed(input=my_tensor)
my_tensor.diag_embed()
torch.diag_embed(input=my_tensor, offset=0, dim1=0, dim2=1)
torch.diag_embed(input=my_tensor, offset=0, dim1=0, dim2=-1)
torch.diag_embed(input=my_tensor, offset=0, dim1=1, dim2=0)
torch.diag_embed(input=my_tensor, offset=0, dim1=1, dim2=-2)
torch.diag_embed(input=my_tensor, offset=0, dim1=-1, dim2=0)
torch.diag_embed(input=my_tensor, offset=0, dim1=-1, dim2=-2)
torch.diag_embed(input=my_tensor, offset=0, dim1=-2, dim2=1)
torch.diag_embed(input=my_tensor, offset=0, dim1=-2, dim2=-1)
# tensor([[7, 0, 0],
#         [0, -4, 0],
#         [0, 0, 5]])

torch.diag_embed(input=my_tensor, offset=1, dim1=0, dim2=1)
torch.diag_embed(input=my_tensor, offset=1, dim1=0, dim2=-1)
torch.diag_embed(input=my_tensor, offset=1, dim1=-2, dim2=1)
torch.diag_embed(input=my_tensor, offset=1, dim1=-2, dim2=-1)
torch.diag_embed(input=my_tensor, offset=-1, dim1=1, dim2=0)
torch.diag_embed(input=my_tensor, offset=-1, dim1=1, dim2=-2)
torch.diag_embed(input=my_tensor, offset=-1, dim1=-1, dim2=0)
torch.diag_embed(input=my_tensor, offset=-1, dim1=-1, dim2=-2)
# tensor([[0, 7, 0, 0],
#         [0, 0, -4, 0],
#         [0, 0, 0, 5],
#         [0, 0, 0, 0]])

torch.diag_embed(input=my_tensor, offset=1, dim1=1, dim2=0)
torch.diag_embed(input=my_tensor, offset=1, dim1=1, dim2=-2)
torch.diag_embed(input=my_tensor, offset=1, dim1=-1, dim2=0)
torch.diag_embed(input=my_tensor, offset=1, dim1=-1, dim2=-2)
torch.diag_embed(input=my_tensor, offset=-1, dim1=0, dim2=1)
torch.diag_embed(input=my_tensor, offset=-1, dim1=0, dim2=-1)
torch.diag_embed(input=my_tensor, offset=-1, dim1=-2, dim2=1)
torch.diag_embed(input=my_tensor, offset=-1, dim1=-2, dim2=-1)
# tensor([[0, 0, 0, 0],
#         [7, 0, 0, 0],
#         [0, -4, 0, 0],
#         [0, 0, 5, 0]])

torch.diag_embed(input=my_tensor, offset=2, dim1=0, dim2=1)
torch.diag_embed(input=my_tensor, offset=2, dim1=0, dim2=-1)
torch.diag_embed(input=my_tensor, offset=2, dim1=-2, dim2=1)
torch.diag_embed(input=my_tensor, offset=2, dim1=-2, dim2=-1)
torch.diag_embed(input=my_tensor, offset=-2, dim1=1, dim2=0)
torch.diag_embed(input=my_tensor, offset=-2, dim1=1, dim2=-2)
torch.diag_embed(input=my_tensor, offset=-2, dim1=-1, dim2=0)
torch.diag_embed(input=my_tensor, offset=-2, dim1=-1, dim2=-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_embed(input=my_tensor, offset=2, dim1=1, dim2=0)
torch.diag_embed(input=my_tensor, offset=2, dim1=1, dim2=-2)
torch.diag_embed(input=my_tensor, offset=2, dim1=-1, dim2=0)
torch.diag_embed(input=my_tensor, offset=2, dim1=-1, dim2=-2)
torch.diag_embed(input=my_tensor, offset=-2, dim1=0, dim2=1)
torch.diag_embed(input=my_tensor, offset=-2, dim1=0, dim2=-1)
torch.diag_embed(input=my_tensor, offset=-2, dim1=-2, dim2=1)
torch.diag_embed(input=my_tensor, offset=-2, dim1=-2, dim2=-1)
# 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_embed(input=my_tensor)
# tensor([[7., 0., 0.],
#         [0., -4., 0.],
#         [0., 0., 5.]])

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

torch.diag_embed(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_embed(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_embed(input=my_tensor)
torch.diag_embed(input=my_tensor, offset=0, dim1=0, dim2=1)
torch.diag_embed(input=my_tensor, offset=0, dim1=0, dim2=-2)
torch.diag_embed(input=my_tensor, offset=0, dim1=1, dim2=0)
torch.diag_embed(input=my_tensor, offset=0, dim1=1, dim2=-3)
torch.diag_embed(input=my_tensor, offset=0, dim1=-2, dim2=0)
torch.diag_embed(input=my_tensor, offset=0, dim1=-2, dim2=-3)
torch.diag_embed(input=my_tensor, offset=0, dim1=-3, dim2=1)
torch.diag_embed(input=my_tensor, offset=0, dim1=-3, dim2=-2)
# tensor([[[7, -6, 9],
#          [0, 0, 0],
#          [0, 0, 0]],
#         [[0, 0, 0],
#          [-4, -3, 1],
#          [0, 0, 0]],
#         [[0, 0, 0],
#          [0, 0, 0],
#          [5, 8, -2]]])

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

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

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

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

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

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

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

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

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

torch.diag_embed(input=my_tensor, offset=2, dim1=0, dim2=2)
torch.diag_embed(input=my_tensor, offset=2, dim1=0, dim2=-1)
torch.diag_embed(input=my_tensor, offset=2, dim1=-3, dim2=2)
torch.diag_embed(input=my_tensor, offset=2, dim1=-3, dim2=-1)
torch.diag_embed(input=my_tensor, offset=-2, dim1=2, dim2=0)
torch.diag_embed(input=my_tensor, offset=-2, dim1=2, dim2=-3)
torch.diag_embed(input=my_tensor, offset=-2, dim1=-1, dim2=0)
torch.diag_embed(input=my_tensor, offset=-2, dim1=-1, dim2=-3)
# tensor([[[0, 0, 7, 0, 0],
#          [0, 0, -6, 0, 0],
#          [0, 0, 9, 0, 0]],
#         [[0, 0, 0, -4, 0],
#          [0, 0, 0, -3, 0],
#          [0, 0, 0, 1, 0]],
#         [[0, 0, 0, 0, 5],
#          [0, 0, 0, 0, 8],
#          [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, 0, 0, 0],
#          [0, 0, 0, 0, 0]]])

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

torch.diag_embed(input=my_tensor, offset=2, dim1=1, dim2=2)
torch.diag_embed(input=my_tensor, offset=2, dim1=1, dim2=-1)
torch.diag_embed(input=my_tensor, offset=2, dim1=-2, dim2=2)
torch.diag_embed(input=my_tensor, offset=2, dim1=-2, dim2=-1)
torch.diag_embed(input=my_tensor, offset=-2, dim1=2, dim2=1)
torch.diag_embed(input=my_tensor, offset=-2, dim1=2, dim2=-2)
torch.diag_embed(input=my_tensor, offset=-2, dim1=-1, dim2=1)
torch.diag_embed(input=my_tensor, offset=-2, dim1=-1, dim2=-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]],
#         [[0, 0, -6, 0, 0],
#          [0, 0, 0,-3, 0],
#          [0, 0, 0, 0, 8],
#          [0, 0, 0, 0, 0],
#          [0, 0, 0, 0, 0]],
#         [[0, 0, 9, 0, 0],
#          [0, 0, 0, 1, 0],
#          [0, 0, 0, 0, -2],
#          [0, 0, 0, 0, 0],
#          [0, 0, 0, 0, 0]]])

torch.diag_embed(input=my_tensor, offset=2, dim1=2, dim2=0)
torch.diag_embed(input=my_tensor, offset=2, dim1=2, dim2=-3)
torch.diag_embed(input=my_tensor, offset=2, dim1=-1, dim2=0)
torch.diag_embed(input=my_tensor, offset=2, dim1=-1, dim2=-3)
torch.diag_embed(input=my_tensor, offset=-2, dim1=0, dim2=2)
torch.diag_embed(input=my_tensor, offset=-2, dim1=0, dim2=-1)
torch.diag_embed(input=my_tensor, offset=-2, dim1=-3, dim2=2)
torch.diag_embed(input=my_tensor, offset=-2, dim1=-3, dim2=-1)
# tensor([[[0, 0, 0, 0, 0],
#          [0, 0, 0, 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],
#          [-6, 0, 0, 0, 0],
#          [9, 0, 0, 0, 0]],
#         [[0, -4, 0, 0, 0],
#          [0, -3, 0, 0, 0],
#          [0, 1, 0, 0, 0]],
#         [[0, 0, 5, 0, 0],
#          [0, 0, 8, 0, 0],
#          [0, 0, -2, 0, 0]]])

torch.diag_embed(input=my_tensor, offset=2, dim1=2, dim2=1)
torch.diag_embed(input=my_tensor, offset=2, dim1=2, dim2=-2)
torch.diag_embed(input=my_tensor, offset=2, dim1=-1, dim2=1)
torch.diag_embed(input=my_tensor, offset=2, dim1=-1, dim2=-2)
torch.diag_embed(input=my_tensor, offset=-2, dim1=1, dim2=2)
torch.diag_embed(input=my_tensor, offset=-2, dim1=1, dim2=-1)
torch.diag_embed(input=my_tensor, offset=-2, dim1=-2, dim2=2)
torch.diag_embed(input=my_tensor, offset=-2, dim1=-2, dim2=-1)
# 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]],
#         [[0, 0, 0, 0, 0],
#          [0, 0, 0, 0, 0],
#          [-6, 0, 0, 0, 0],
#          [0, -3, 0, 0, 0],
#          [0, 0, 8, 0, 0]],
#         [[0, 0, 0, 0, 0],
#          [0, 0, 0, 0, 0],
#          [9, 0, 0, 0, 0],
#          [0, 1, 0, 0, 0],
#          [0, 0, -2, 0, 0]]])
```