Custom discrete distribution on GPU

I am trying to generate a random matrix that is filled with a discrete distribution with three values. In my case, my distribution looks like -1 with probability 1/4, +1 with probability -1/4, and 0 with probability 1/2. Is there an elegant way to do this in Pytorch? I am working with the C++ frontend, but would be willing to switch to the python frontend to get this working.


I do not know if my solution is elegant, but I believe it is a way to do what you want to do!

    x = torch.tensor([-1, 0, 0, 1], dtype=torch.float32)
    y = torch.zeros(size=(10, 10), dtype=torch.float32)
    for i in range(y.shape[0]):
        for j in range(y.shape[1]):
            choice = torch.randint(high=x.shape[0], size=(1, ))
            y[i, j] = x[choice]
    print(torch.sum((y == 1.0)))
    print(torch.sum((y == 0.0)))
    print(torch.sum((y == -1.0)))

Good luck!