ashered
ashered

Reputation: 79

How do I make a mask of diagonal matrix, but starting from the 2nd column?

So here is what I can get with torch.eye(3,4) now

The matrix I get:

[[1, 0, 0, 0],
 [0, 1, 0, 0],
 [0, 0, 1, 0]]

Is there any (easy)way to transform it, or make such a mask in this format:

The matrix I want:

[[0, 1, 0, 0],
 [0, 0, 1, 0],
 [0, 0, 0, 1]]

Upvotes: 3

Views: 2966

Answers (2)

kmario23
kmario23

Reputation: 61365

Here is another solution using torch.diagflat(), and using a positive offset for shifting/moving the diagonal above the main diagonal.

# diagonal values to fill
In [253]: diagonal_vals = torch.ones(3, dtype=torch.long)  

# desired tensor but ...
In [254]: torch.diagflat(diagonal_vals, offset=1) 
Out[254]: 
tensor([[0, 1, 0, 0],
        [0, 0, 1, 0],
        [0, 0, 0, 1],
        [0, 0, 0, 0]])

The above operation gives us a square matrix; however, we need a non-square matrix of shape (3,4). So, we'll just ignore the last row with simple indexing:

# shape (3, 4) with 1's above the main diagonal
In [255]: torch.diagflat(diagonal_vals, offset=1)[:-1] 
Out[255]: 
tensor([[0, 1, 0, 0],
        [0, 0, 1, 0],
        [0, 0, 0, 1]])

Upvotes: 1

kuzand
kuzand

Reputation: 9806

You can do it by using torch.diagonal and specifying the diagonal you want:

>>> torch.diag(torch.tensor([1,1,1]), diagonal=1)[:-1]

tensor([[0, 1, 0, 0],
        [0, 0, 1, 0],
        [0, 0, 0, 1]])
  • If :attr:diagonal = 0, it is the main diagonal.
  • If :attr:diagonal > 0, it is above the main diagonal.
  • If :attr:diagonal < 0, it is below the main diagonal.

Upvotes: 2

Related Questions