CODEWITHSUNDEEP
pythonpytorchvectorizationtensor
Ethan
Ethan

Reputation: 387

2-D Vectorized Function with pytorch

I have a function that I need to compute at every entry of a 2-D torch tensor, and it is depends on the value of the index of both axes. Right now, I have only been able to implement this as a nested for loop, iterating over both axes. This is quite slow (and needs to be executed >10^5 times) and I want to speed it up for better scaling.

vs = 200
nt = 12
b = torch.ones(vs)/vs
n_kw = torch.rand((nt, vs))
n_k = torch.rand((nt,))

def estimate_p(nt, vs, n_kw, n_k):
    p = torch.zeros((nt, vs))
    
    for i in range(0, nt):
        for j in range(0, vs):
            p[i,j] = (n_kw[i, j] + b[j])/(n_k[i] + torch.sum(b))
    return p

Is there a way to vectorize this/map based in the i,j index?

Upvotes: 1

Views: 140

Answers (1)

Quang Hoang
Quang Hoang

Reputation: 150735

Try playing with broadcasting:

def estimate_p(nt, vs, n_kw, n_k):
    return (n_kw + b) / (n_k + b.sum()).unsqueeze(-1)

Upvotes: 1

Related Questions