Fonction torch.diag_embed de PyTorch

Pytorch torch 参考手册Manuel de référence Pytorch torch

torch.diag_embedest une fonction de PyTorch utilisée pour intégrer le tenseur d’entrée comme éléments diagonaux dans un nouveau tenseur. Elle crée une diagonale sur la dernière dimension de l’entrée.

Définition de fonction

torch.diag_embed(input, diagonal=0, offset=0)

Description des paramètres :

  • input: tenseur d'entrée
  • diagonal/offset: indice diagonal

Exemples d'utilisation

Exemple

import torch

# Créer un tenseur unidimensionnel
x = torch.tensor([1, 2, 3])

# Créer un plongement diagonal
y = torch.diag_embed(x)
print(y)

La sortie est :

tensor([[1, 0, 0],
        [0, 2, 0],
        [0, 0, 3]])

Exemple

import torch

# Créer un tenseur bidimensionnel
x = torch.tensor([[1, 2], [3, 4]])

# Créer un plongement diagonal
y = torch.diag_embed(x)
print(y.shape)
print(y)

La sortie est :

torch.Size([2, 2, 2])
tensor([[[1, 0],
         [0, 2]],

        [[3, 0],
         [0, 4]]])

Pytorch torch 参考手册Manuel de référence Pytorch torch

Autres extensions