Flatten upper triangular matrix

Given an NxN matrix Mat, I would like to have a vector of upper diagonal elements Vec of the matrix. Are there anyways, without using a nested loop, to do this? With a nested loop

Vec = None
for i in range(N):
    for j in range(N):
        if i>j:
            if Vec is None:
                Vec = adjacency[:, i, j].unsqueeze(1)
                Vec = torch.cat((Vec, adjacency[:, i, j].unsqueeze(1)), dim=1)

where Vec is an N(N-1)/2-dimensional vector.


You should be able to use triu_indices() here to get the indices and then just index adjacency with it using a single function call.

1 Like