Dear experts,

I am trying to find the best way to perform a series of weighted random sampling, each from a different distribution. I will try to elaborate on what I intend to do (Maybe if you just jump to the code it will be immediately clear).

I have an input `x`

of length `L`

. Each element is an integer between 0 and `N`

.

I want to generate an output `y`

of length `L`

, where each element of y is an integer between 0 and `M`

.

Each element `y[i]`

is transitioned from the corresponding element `x[i]`

, with a certain transition probability distribution defined by a `N`

x`M`

stochastic matrix `p`

.

`p[i][j]`

denotes the probability of `y`

becoming `j`

, if `x`

is `i`

.

Basically, if `N`

==`M`

, I am trying to do a Markov chain transition.

Below is a (very slow) code I implemented.

```
import torch
from torch import nn
N = 10
M = 20
L = 5
p = torch.rand([N, M])
x = torch.randint(0, N, [L])
print(x)
y = torch.zeros([L])
for i in range(L):
y[i] = list(torch.utils.data.WeightedRandomSampler(p[x[i]], 1))[0]
print(y)
```

Here, I am not normalizing `p`

row-wise (i.e., it is not a stochastic matrix) because `WeightedRandomSampler`

does not need it to be normalized row-wise.

The problem with this code is that I am using a `for`

loop, and this becomes too slow for large `L`

.

I wonder if there is a way to implement the same thing faster, without having the `for`

loop.

Any insights would be highly appreciated.

Thank you.