Pytoch: quickly find the main diagonal elements and non diagonal elements of NxN matrix

Antecedents

In the process of writing code today , To get one respectively N x NTensor Main diagonal elements and off diagonal elements of matrix , See a clever way of writing , So write it down .

Explore

torch.diagonal()

The main diagonal element is easy to get ,Pytorch Have a ready-made API Callable , by torch.diagonal, Details are as follows ：

It's also easy to use , Examples are as follows ：

x = torch.randn(4,4)
# tensor([[ 0.9148, 0.1396, -0.8974, 2.0014],
# [ 0.1129, -0.3656, 0.4371, 0.2618],
# [ 1.1049, -0.0774, -0.4160, -0.4922],
# [ 1.3197, -0.2022, -0.0031, -1.3811]])

torch.diagonal(x)
# tensor([ 0.9148, -0.3656, -0.4160, -1.3811])

Clever use of matrix transformation

About non diagonal elements , There is no specific API 了 , In the process of searching for information , See a more ingenious method , Go straight to the code ：

Main reference ：https://github1s.com/facebookresearch/barlowtwins/blob/e6f34a01c0cde6f05da6f431ef8a577b42e94e71/main.py#L207

n, m = x.shape
assert n == m
x.flatten()[:-1].view(n-1,n+1)[:,1:].flatten()
# tensor([ 0.1396, -0.8974, 2.0014, 0.1129, 0.4371, 0.2618, 1.1049, -0.0774,
# -0.4922, 1.3197, -0.2022, -0.0031])

The core code is on the last line , The following is mainly about decomposition . The first use of flatten() Straighten vector , Then remove the last element , obtain $n^2-1$ Elements , Then construct a dimension of [N-1, N+1] Matrix . In this matrix , All previous diagonal elements appear in the 1 Column , As shown below ：

Then get according to the index [:, 1:] Elements , The result is the non diagonal element of the original matrix .