Equivalent of tf.linalg.diag_part in PyTorch
As I'm reimplementing some code, I'm wondering if there is any equivalent of tf.linalg.diag_part (docs) in PyTorch ..?

I don't believe there's a direct equivalent. However, you can get away using torch.diag:
&gt;&gt;&gt; x = torch.tensor([[1, 2, 3, 4],
                      [5, 6, 7, 8]])

&gt;&gt;&gt; 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]]])

