torch.diag¶
- torch.diag(input, diagonal=0, *, out=None) Tensor ¶
如果
input
是一个向量 (1-D 张量),则返回一个 2-D 方阵张量,其对角线元素由input
的元素构成。如果
input
是一个矩阵 (2-D 张量),则返回一个 1-D 张量,其元素为input
的对角线元素。
参数
diagonal
控制考虑哪条对角线示例
获取以输入向量为对角线的方阵
>>> a = torch.randn(3) >>> a tensor([ 0.5950,-0.0872, 2.3298]) >>> torch.diag(a) tensor([[ 0.5950, 0.0000, 0.0000], [ 0.0000,-0.0872, 0.0000], [ 0.0000, 0.0000, 2.3298]]) >>> torch.diag(a, 1) tensor([[ 0.0000, 0.5950, 0.0000, 0.0000], [ 0.0000, 0.0000,-0.0872, 0.0000], [ 0.0000, 0.0000, 0.0000, 2.3298], [ 0.0000, 0.0000, 0.0000, 0.0000]])
获取给定矩阵的第 k 条对角线
>>> a = torch.randn(3, 3) >>> a tensor([[-0.4264, 0.0255,-0.1064], [ 0.8795,-0.2429, 0.1374], [ 0.1029,-0.6482,-1.6300]]) >>> torch.diag(a, 0) tensor([-0.4264,-0.2429,-1.6300]) >>> torch.diag(a, 1) tensor([ 0.0255, 0.1374])