PyTorch torch.diag_embed Function
Pytorch torch Reference Manual
torch.diag_embedis a function in PyTorch used to embed the input tensor as diagonal elements into a new tensor. It creates a diagonal on the last dimension of the input.
Function Definition
torch.diag_embed(input, diagonal=0, offset=0)
Parameter Description:
input: input tensordiagonal/offset: diagonal index
Usage Examples
Example
import torch
# Create a one-dimensional tensor
x = torch.tensor([1, 2, 3])
# Create diagonal embedding
y = torch.diag_embed(x)
print(y)
# Create a one-dimensional tensor
x = torch.tensor([1, 2, 3])
# Create diagonal embedding
y = torch.diag_embed(x)
print(y)
The output is:
tensor([[1, 0, 0],
[0, 2, 0],
[0, 0, 3]])
Example
import torch
# Create a two-dimensional tensor
x = torch.tensor([[1, 2], [3, 4]])
# Create diagonal embedding
y = torch.diag_embed(x)
print(y.shape)
print(y)
# Create a two-dimensional tensor
x = torch.tensor([[1, 2], [3, 4]])
# Create diagonal embedding
y = torch.diag_embed(x)
print(y.shape)
print(y)
The output is:
torch.Size([2, 2, 2])
tensor([[[1, 0],
[0, 2]],
[[3, 0],
[0, 4]]])
Other Extensions