PyTorch 中 tf.linalg.diag_part 的等价物

【问题标题】：Equivalent of tf.linalg.diag_part in PyTorchPyTorch 中 tf.linalg.diag_part 的等价物
【发布时间】：2022-01-19 16:56:14
【问题描述】：

当我重新实现一些代码时，我想知道PyTorch 中是否有任何与tf.linalg.diag_part (docs) 等价的东西？

【问题讨论】：

标签： tensorflow pytorch

【解决方案1】：

我不相信有直接的等价物。但是，您可以使用torch.diag 逃脱：

>>> x = torch.tensor([[1, 2, 3, 4],
                      [5, 6, 7, 8]])

>>> torch.diag(x.flatten()).reshape(-1, 4, 2, 4).sum(-2)
tensor([[[1, 0, 0, 0],
         [0, 2, 0, 0],
         [0, 0, 3, 0],
         [0, 0, 0, 4]],

        [[5, 0, 0, 0],
         [0, 6, 0, 0],
         [0, 0, 7, 0],
         [0, 0, 0, 8]]])

【讨论】：