Creating tensor with values depending on the indices

I need to create a tensor where entry i,j has value abs(i - j).
For now, the only solution I’ve found is by cycling each entry with for loop:

A = torch.zeros(1000, 2000)
for i in range(1000):
    for j in range(2000):
        A[i, j] = abs(i - j)

Unfortunately this solution is extremely slow.

I need this tensor to use it as mask of another tensor.

Someone has any alternative idea?

I’m not sure about the exact code but you can use range and expand to have matrices like A[i, j] = i, forall j: range(1000).unsqueeze(1).expand(1000, 2000). And the same for the other one with unsqueeze(0).
Then you can substract them element-wise and take the absolute value of the result.

Really thanks, I was using python list multiplication with torch.stack, but this is much faster.


In [22]: %timeit torch.stack([torch.arange(1000)] * 2000, 1)                                   
42.3 ms ± 125 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [23]: %timeit torch.arange(1000).unsqueeze(1).expand(1000, 2000)                            
13.1 µs ± 21.5 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)


In [27]: %timeit torch.stack([torch.arange(1000).to(device)] * 2000, 1).to(device)             
2.15 ms ± 6.66 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [30]: %timeit torch.arange(1000).to(device).unsqueeze(1).expand(1000, 2000).to(device)    
43.8 µs ± 152 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
1 Like