Create a sliding mask around diagonals

Hello,

Firstly I’m very sorry with this newbie question. I couldn’t find an efficient answer. I want to create a mask with zeros and ones for the n by n matrix. For value each i-th row I want to set the columns between i-a and i+a to 1 and rest of them to 0.
As an example of the values
n = 4 (4x4 matrix)
a = 1
I want to get a matrix as follows.
[
[1,1,0,0],
[1,1,1,0],
[0,1,1,1],
[0,0,1,1]
]

How can I implement this in the best way in Pytorch? Thanks in advance…

This piece of code can do the job. But it’s not very optimal.

import torch

n = 4
a = 1

A = torch.zeros((n, n))

for i in range(n) :
    for j in range(max(i - a, 0), min(i + a + 1, n)) :
        A[i][j] = 1.0

Thank you I am using in this way but I want to do it without a loop.

Hi Talip!

I think that torch.tril() and/or torch.triu() will be the key to what you want.

Here is one specific approach:

>>> torch.__version__
'1.6.0'
>>> n = 4
>>> a = 1
>>> torch.tril (torch.ones ((n, n)), a) * torch.triu (torch.ones ((n, n)), -a)
tensor([[1., 1., 0., 0.],
        [1., 1., 1., 0.],
        [0., 1., 1., 1.],
        [0., 0., 1., 1.]])

Best.

K. Frank

thank you very much Frank!