Differences with torch.diag

The following mapping relationships can be found in this file.

PyTorch APIs	MindSpore APIs
torch.diag	mindspore.ops.diag
torch.Tensor.diag	mindspore.Tensor.diag

torch.diag

torch.diag(input, diagonal=0, *, out=None) -> Tensor

For more information, see torch.diag.

mindspore.ops.diag

mindspore.ops.diag(input) -> Tensor

For more information, see mindspore.ops.diag.

Differences

API function of MindSpore is not consistent with that of PyTorch.

PyTorch: If the input is a one-dimensional tensor, the diagonal tensor is constructed by using the one-dimensional tensor composed of the diagonal values of the input. If the input is a matrix, the one-dimensional tensor composed of the diagonal elements of the input is returned.

MindSpore: If the input is a one-dimensional tensor, MindSpore API achieves the same function as PyTorch. If the input is a matrix, it does not achieve the same function as PyTorch, and there is no diagonal parameter to control the position of the diagonals.

Categories	Subcategories	PyTorch	MindSpore	Difference
Parameters	Parameter 1	input	input	-
	Parameter 2	diagonal	-	The value of `diagonal` in PyTorch is used to control the position of the diagonal. MindSpore does not have this parameter.
	Parameter 3	out	-	For details, see General Difference Parameter Table

Code Example 1

The PyTorch API parameter x supports both multidimensional and one-dimensional tensors, and there is a diagonal parameter to control the position of the diagonal, while the MindSpore API does not have a diagonal parameter. When the input parameter x is a one-dimensional tensor and diagonal is 0, the two APIs achieve the same function.

# PyTorch
import torch
x = torch.tensor([1,2,3,4],dtype=int)
out = torch.diag(x)
out = out.detach().numpy()
print(out)
# [[1 0 0 0]
#  [0 2 0 0]
#  [0 0 3 0]
#  [0 0 0 4]]

# MindSpore
from mindspore import Tensor
import mindspore.ops as ops
input_x = Tensor([1, 2, 3, 4]).astype('int32')
output = ops.diag(input_x)
print(output)
# [[1 0 0 0]
#  [0 2 0 0]
#  [0 0 3 0]
#  [0 0 0 4]]

Code Example 2

When the input parameter x is a one-dimensional tensor and diagonal is not 0, this API of PyTorch controls the position of the diagonal, while this API of MindSpore does not have a diagonal parameter, and the output obtained from this API can be processed by mindspore.ops.pad to achieve the same function.

# PyTorch
import torch
x = torch.tensor([1,2,3,4],dtype=int)
# Results for diagonal greater than 0
out = torch.diag(x, diagonal=1)
out = out.detach().numpy()
print(out)
# [[0 1 0 0 0]
#  [0 0 2 0 0]
#  [0 0 0 3 0]
#  [0 0 0 0 4]
#  [0 0 0 0 0]]

# Results for diagonal smaller than 0
out = torch.diag(x, diagonal=-1)
out = out.detach().numpy()
print(out)
# [[0 0 0 0 0]
#  [1 0 0 0 0]
#  [0 2 0 0 0]
#  [0 0 3 0 0]
#  [0 0 0 4 0]]

# MindSpore
from mindspore import Tensor
import mindspore.ops as ops
input_x = Tensor([1, 2, 3, 4]).astype('int32')
output = ops.diag(input_x)
# MindSpore implements this API function when the diagonal is greater than 0.
padding = ((1, 0, 0, 1))
a = ops.pad(output, padding)
print(a)
# [[0 1 0 0 0]
#  [0 0 2 0 0]
#  [0 0 0 3 0]
#  [0 0 0 0 4]
#  [0 0 0 0 0]]

# MindSpore implements this API function when the diagonal is smaller than 0.
padding = ((0, 1, 1, 0))
a = ops.pad(output, padding)
print(a)
# [[0 0 0 0 0]
#  [1 0 0 0 0]
#  [0 2 0 0 0]
#  [0 0 3 0 0]
#  [0 0 0 4 0]]

Code Example 3

This API of PyTorch is used to extract a one-dimensional tensor composed of diagonals when the input is a matrix and diagonal is used, MindSpore does not support this function. Using mindspore.numpy.diag can implement this function.

# PyTorch
import torch
x = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]],dtype=int)
# Results for diagonal greater than 0
out = torch.diag(x, diagonal=1)
out = out.detach().numpy()
print(out)
# [2 6]

# Result when diagonal is the default value 0
out = torch.diag(x)
out = out.detach().numpy()
print(out)
# [1 5 9]

# Results for diagonal smaller than 0
out = torch.diag(x, diagonal=-1)
out = out.detach().numpy()
print(out)
# [4 8]

# MindSpore
from mindspore import Tensor
import mindspore.ops as ops
import mindspore.numpy as np
input_x = Tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]]).astype('int32')
# This function of mindspore.numpy.diag when diagonal is greater than 0
output = np.diag(input_x, k=1)
print(output)
# [2 6]

# This function of mindspore.numpy.diag when diagonal defaults to 0
output = np.diag(input_x)
print(output)
# [1 5 9]

# This function of mindspore.numpy.diag when diagonal is smaller than 0
output = np.diag(input_x, k=-1)
print(output)
# [4 8]