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 &gt;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?

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

