# How to produce these tensors efficiently/fast?

I would like to produce the following tensor of size (N*N) where the ones (D) appear as follows:

create_mask(N=10, D=5)

tensor([[1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
[1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
[1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
[0, 1, 1, 1, 1, 1, 0, 0, 0, 0],
[0, 0, 1, 1, 1, 1, 1, 0, 0, 0],
[0, 0, 0, 1, 1, 1, 1, 1, 0, 0],
[0, 0, 0, 0, 1, 1, 1, 1, 1, 0],
[0, 0, 0, 0, 0, 1, 1, 1, 1, 1],
[0, 0, 0, 0, 0, 1, 1, 1, 1, 1],
[0, 0, 0, 0, 0, 1, 1, 1, 1, 1]])


or

create_mask(N=5, D=3)

tensor([[1, 1, 1, 0, 0],
[1, 1, 1, 0, 0],
[0, 1, 1, 1, 0],
[0, 0, 1, 1, 1],
[0, 0, 1, 1, 1]])


I need this function to be very efficient/fast because I will be generating large tensors (usually N~[4000,10000]).

Here is my current approach (too slow) :

import torch
import numpy as np

assert num_around & 1
num_neighbor = num_around//2

output = np.array([np.array(range(i-num_neighbor, i+num_neighbor+1)) for i in range(num_words)])
output[:num_neighbor, :] = output[num_neighbor]
output[-num_neighbor:, :] = output[-num_neighbor-1]
output = torch.tensor(output)