Replacing this line by
```
            p_w_sample = p_w_c.take(y, axis=1).sum(axis=1)
```
lead to a significant speed up.