# Sampling from a probability distribution on GPU

**URL:** <https://discourse.julialang.org/t/sampling-from-a-probability-distribution-on-gpu/94389>\
**Category:** GPU\
**Tags:** statistics, gpu, distributions\
**Created:** [February 10, 2023, 2:01am UTC](https://discourse.julialang.org/t/sampling-from-a-probability-distribution-on-gpu/94389 "2023-02-10T02:01:06Z")\
**Posts on this page:** 14\
**Page:** 1

<div class="post-metadata">

**Author:** ![ArjunNarayanan](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/arjunnarayanan/32/6475_2.png) [@ArjunNarayanan](https://discourse.julialang.org/u/ArjunNarayanan)\
**Post date:** [February 10, 2023, 2:01am UTC](https://discourse.julialang.org/t/sampling-from-a-probability-distribution-on-gpu/94389/1 "2023-02-10T02:01:06Z")

</div>

I’m interested in doing something like

```julia
using Distributions: Categorical
using CUDA

p = cu([0.7,0.2,0.1])
idx = rand(Categorical(p))

```

Is there a way to sample from a vector of probability masses on the GPU? This kind of stuff comes up a fair amount in machine learning.

---

<div class="post-metadata">

**Author:** ![rmsmsgood](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/rmsmsgood/32/20544_2.png) [@rmsmsgood](https://discourse.julialang.org/u/rmsmsgood)\
**Post date:** [February 10, 2023, 2:09am UTC](https://discourse.julialang.org/t/sampling-from-a-probability-distribution-on-gpu/94389/2 "2023-02-10T02:09:35Z")

</div>

You mean, easy or convenient way? In native CUDA? If not, I guess that your job is not much difficult to implement.

---

<div class="post-metadata">

**Author:** ![ArjunNarayanan](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/arjunnarayanan/32/6475_2.png) [@ArjunNarayanan](https://discourse.julialang.org/u/ArjunNarayanan)\
**Post date:** [February 10, 2023, 2:18am UTC](https://discourse.julialang.org/t/sampling-from-a-probability-distribution-on-gpu/94389/3 "2023-02-10T02:18:25Z")

</div>

I guess something that’s convenient? I don’t have much experience with sampling algorithms. But if something needs to be implemented, I could do it with the right guidance.

---

<div class="post-metadata">

**Author:** ![jpsamaroo](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/jpsamaroo/32/46804_2.png) [@jpsamaroo](https://discourse.julialang.org/u/jpsamaroo)\
**Post date:** [February 10, 2023, 11:30pm UTC](https://discourse.julialang.org/t/sampling-from-a-probability-distribution-on-gpu/94389/4 "2023-02-10T23:30:21Z")

</div>

Is your intention to be able to call `rand(Categorical(p))` from within a GPU kernel? Or does the code you posted just not work, and you want to know how to make it work?

---

<div class="post-metadata">

**Author:** ![jpsamaroo](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/jpsamaroo/32/46804_2.png) [@jpsamaroo](https://discourse.julialang.org/u/jpsamaroo)\
**Post date:** [February 10, 2023, 11:40pm UTC](https://discourse.julialang.org/t/sampling-from-a-probability-distribution-on-gpu/94389/5 "2023-02-10T23:40:01Z")

</div>

This is not a particularly helpful answer; if it’s not difficult to implement, then maybe you could instead point to some resources that would help the OP?

---

<div class="post-metadata">

**Author:** ![simsurace](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/simsurace/32/30216_2.png) [@simsurace](https://discourse.julialang.org/u/simsurace)\
**Post date:** [February 11, 2023, 1:13am UTC](https://discourse.julialang.org/t/sampling-from-a-probability-distribution-on-gpu/94389/6 "2023-02-11T01:13:42Z")

</div>

This type of question comes up often.  
Just a few days ago I saw and answered a very similar question on Slack.

I happen to have [some CUDA.jl kernel code](https://github.com/Theoretical-Neuroscience-Group/BinomialSynapses.jl/blob/c2101b0af2942e7f99c5dc1b1257777818279ddd/src/resample.jl#L65-L156) in a package of mine that I haven’t touched in a couple of years that could serve as a starting point. It is using a naive algorithm though.

```julia
using CUDA
using BinomialSynapses: indices!

function rand_categorical(p, n)
     v = repeat(p', n ÷ length(p) + 1, 1)
     idx = last(indices!(v))
     return idx[1:n]
end

p = cu([0.7,0.2,0.1]) # does not need to be normalized
samples = rand_categorical(p, 1000)

```

should give you a `CuVector` of length 1000 of categorical samples.

But since that function was written for a specific application (resampling particles) where the number of samples needed was equal to the length of `p`, and where there were a lot of different `p`s, this is not going to be competitive in performance if you have short `p` and need lots of samples. You are (much) better off just copying your `p` to the CPU and calling `rand` there, and then copying the samples back to the GPU if needed.

An algorithm that is efficient for lots of repeated samples is probably going to use alias tables. See e.g. [GitHub - ByteHamster/alias-table-gpu: Efficient construction of and sampling from alias tables on the GPU](https://github.com/ByteHamster/alias-table-gpu) and the associated paper.

It would be nice to have a library for efficient/state-of-the-art sampling algorithms on GPUs using some portable approach like KernelAbstractions.jl?

---

<div class="post-metadata">

**Author:** ![ArjunNarayanan](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/arjunnarayanan/32/6475_2.png) [@ArjunNarayanan](https://discourse.julialang.org/u/ArjunNarayanan)\
**Post date:** [February 13, 2023, 6:01pm UTC](https://discourse.julialang.org/t/sampling-from-a-probability-distribution-on-gpu/94389/7 "2023-02-13T18:01:39Z")

</div>

The latter. I don’t think `Categorical` works with `CuArray`.

---

<div class="post-metadata">

**Author:** ![ArjunNarayanan](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/arjunnarayanan/32/6475_2.png) [@ArjunNarayanan](https://discourse.julialang.org/u/ArjunNarayanan)\
**Post date:** [February 13, 2023, 6:03pm UTC](https://discourse.julialang.org/t/sampling-from-a-probability-distribution-on-gpu/94389/8 "2023-02-13T18:03:28Z")

</div>

Thanks that’s helpful.

My application is in Reinforcement Learning. I only need one sample from my distribution which is the action I will take in the next time step. In this situation, perhaps I’m better off just moving the array to CPU and calling `rand` there?

---

<div class="post-metadata">

**Author:** ![ArjunNarayanan](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/arjunnarayanan/32/6475_2.png) [@ArjunNarayanan](https://discourse.julialang.org/u/ArjunNarayanan)\
**Post date:** [February 13, 2023, 6:07pm UTC](https://discourse.julialang.org/t/sampling-from-a-probability-distribution-on-gpu/94389/9 "2023-02-13T18:07:06Z")

</div>

I looked around online and I guess it’s possible to implement sampling using the inverse transform method for reasonable distributions? I think `CUDA` already provides a `cumsum` for the CDF. So one might only need to implement a binary search on `CuArray` to use this method?

---

<div class="post-metadata">

**Author:** ![simsurace](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/simsurace/32/30216_2.png) [@simsurace](https://discourse.julialang.org/u/simsurace)\
**Post date:** [February 13, 2023, 9:36pm UTC](https://discourse.julialang.org/t/sampling-from-a-probability-distribution-on-gpu/94389/10 "2023-02-13T21:36:09Z")

</div>

Yeah, unless you are running many agents in parallel you don’t need to sample on the GPU.

---

<div class="post-metadata">

**Author:** ![e3c6](https://avatars.discourse-cdn.com/v4/letter/e/e79b87/32.png) [@e3c6](https://discourse.julialang.org/u/e3c6)\
**Post date:** [June 13, 2023, 12:24pm UTC](https://discourse.julialang.org/t/sampling-from-a-probability-distribution-on-gpu/94389/11 "2023-06-13T12:24:35Z")

</div>

I was also wondering what’s the best solution to this.

Another possibility is the Gumbel trick, in which case one only needs a CUDA function that returns the index of the maximum entry in an array.

---

<div class="post-metadata">

**Author:** ![findmyway](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/findmyway/32/4946_2.png) [@findmyway](https://discourse.julialang.org/u/findmyway)\
**Post date:** [June 13, 2023, 12:47pm UTC](https://discourse.julialang.org/t/sampling-from-a-probability-distribution-on-gpu/94389/12 "2023-06-13T12:47:31Z")

</div>

I wrote a blog on it several years ago, not sure if it still works. But should be a good starting point 😉

[https://tianjun.me/essays/Categorical\_Sampling\_on\_GPU\_with\_Julia/](https://tianjun.me/essays/Categorical_Sampling_on_GPU_with_Julia/)

---

<div class="post-metadata">

**Author:** ![e3c6](https://avatars.discourse-cdn.com/v4/letter/e/e79b87/32.png) [@e3c6](https://discourse.julialang.org/u/e3c6)\
**Post date:** [June 13, 2023, 12:49pm UTC](https://discourse.julialang.org/t/sampling-from-a-probability-distribution-on-gpu/94389/13 "2023-06-13T12:49:32Z")

</div>

Thanks. I see you use an alias table approach?

Unfortunately I am using the categorical sampling during training so parameters are changing quickly and maintaining an alias table seems like not the right approach because I won’t generate that many samples at fixed parameter values. But I could be wrong.

---

<div class="post-metadata">

**Author:** ![findmyway](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/findmyway/32/4946_2.png) [@findmyway](https://discourse.julialang.org/u/findmyway)\
**Post date:** [June 13, 2023, 12:52pm UTC](https://discourse.julialang.org/t/sampling-from-a-probability-distribution-on-gpu/94389/14 "2023-06-13T12:52:00Z")

</div>

In that case, I’d prefer the Gumbel trick. The extra allocation is trivial.