Función PyTorch torch.diag


Pytorch torch 参考手册Manual de referencia de Pytorch torch

torch.diagEs una función en PyTorch utilizada para crear una matriz diagonal o extraer los elementos diagonales de un tensor.

Definición de la función

torch.diag(input, diagonal=0, out=None)

Descripción de parámetros:

  • input: tensor de entrada
  • diagonal: índice de la diagonal, 0 indica la diagonal principal, los valores positivos indican la diagonal superior y los valores negativos indican la diagonal inferior

Ejemplo de uso

Ejemplo

import torch

# Crear un tensor unidimensional
x = torch.tensor([1, 2, 3])

# Crear una matriz diagonal
y = torch.diag(x)
print(y)

El resultado de salida es:

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

Ejemplo

import torch

# Extraer los elementos diagonales de la matriz
x = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])

# Extraer la diagonal principal
y = torch.diag(x)
print(y)

El resultado de salida es:

tensor([1, 5, 9])

Pytorch torch 参考手册Manual de referencia de Pytorch torch

Otras extensiones