# Is allocation inevitable when generating random numbers from a categorical distribution?

**URL:** <https://discourse.julialang.org/t/is-allocation-inevitable-when-generating-random-numbers-from-a-categorical-distribution/101718>\
**Category:** General Usage\
**Tags:** question\
**Created:** [July 18, 2023, 8:39am UTC](https://discourse.julialang.org/t/is-allocation-inevitable-when-generating-random-numbers-from-a-categorical-distribution/101718 "2023-07-18T08:39:47Z")\
**Posts on this page:** 10\
**Page:** 1

<div class="post-metadata">

**Author:** ![taka255](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/taka255/32/51426_2.png) [@taka255](https://discourse.julialang.org/u/taka255)\
**Post date:** [July 18, 2023, 8:39am UTC](https://discourse.julialang.org/t/is-allocation-inevitable-when-generating-random-numbers-from-a-categorical-distribution/101718/1 "2023-07-18T08:39:47Z")

</div>

Hi,

I am trying to do in-place random number generation using `Distributions.jl`. However, for some reason, only the random number generation from the Categorical distribution takes extra computation time due to the extra allocations. Here is the code.

```julia
function loop()
    w = rand(10^5)
    normalize!(w,1)
    tmpI = zeros(Int,10^5)
    tmpF = zeros(Float64,10^5)

    @time for i in 1:10^3
        rand!(Categorical(w), tmpI)
    end

    @time for i in 1:10^3
        rand!(Normal(0,1), tmpF)
    end

    @time for i in 1:10^3
        rand!(Gamma(1,1), tmpF)
    end

    @time for i in 1:10^3
        rand!(InverseGamma(1,1), tmpF)
    end

    @time for i in 1:10^3
        rand!(Uniform(0,1), tmpF)
    end
end

loop()

```

And here is the result.

```julia
  4.153208 seconds (10.00 k allocations: 3.726 GiB, 1.42% gc time)
  0.269260 seconds
  0.444689 seconds
  0.883855 seconds
  0.063641 seconds

```

Is this inevitable?

---

<div class="post-metadata">

**Author:** ![nilshg](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/nilshg/32/2283_2.png) [@nilshg](https://discourse.julialang.org/u/nilshg)\
**Post date:** [July 18, 2023, 9:27am UTC](https://discourse.julialang.org/t/is-allocation-inevitable-when-generating-random-numbers-from-a-categorical-distribution/101718/2 "2023-07-18T09:27:33Z")

</div>

Looks like it, although I can save 20% allocations and 10% timing on my machine by pulling out the creation of the distribution object:

```julia
julia> catdist = Categorical(w);

julia> @btime rand!(Categorical($w), $tmpI);
  3.279 ms (10 allocations: 3.81 MiB)

julia> @btime rand!($catdist, $tmpI);
  2.936 ms (8 allocations: 3.05 MiB)

```

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [July 18, 2023, 10:02am UTC](https://discourse.julialang.org/t/is-allocation-inevitable-when-generating-random-numbers-from-a-categorical-distribution/101718/3 "2023-07-18T10:02:41Z")

</div>

> [@nilshg](#):
>
> `catdist = Categorical(w)`

Try using the constructor without argument checks:

```julia
julia> catdist = Categorical(w; check_args=false)

```

The reason is that `Categorical` is an alias for `DiscreteNonParametric`, and its constructor looks like [this](https://github.com/JuliaStats/Distributions.jl/blob/c8d3e4b52ea2c04f655510af85f93fa3876f258f/src/univariate/discrete/discretenonparametric.jl#L20-L36) at the moment:

```julia
struct DiscreteNonParametric{T<:Real,P<:Real,Ts<:AbstractVector{T},Ps<:AbstractVector{P}} <: DiscreteUnivariateDistribution
    support::Ts
    p::Ps

    function DiscreteNonParametric{T,P,Ts,Ps}(xs::Ts, ps::Ps; check_args::Bool=true) where {
            T<:Real,P<:Real,Ts<:AbstractVector{T},Ps<:AbstractVector{P}}
        check_args || return new{T,P,Ts,Ps}(xs, ps)
        @check_args(
            DiscreteNonParametric,
            (length(xs) == length(ps), "length of support and probability vector must be equal"),
            (ps, isprobvec(ps), "vector is not a probability vector"),
            (xs, allunique(xs), "support must contain only unique elements"),
        )
        sort_order = sortperm(xs)
        new{T,P,Ts,Ps}(xs[sort_order], ps[sort_order])
    end
end

```

Note the presence of `sortperm(xs)` as well as reindexing with `sort_order`.  
I think sorting allows for faster sampling (not sure, cause intuitively I would sort `ps`), but it also means we waste a lot of time during construction as a result.  
Strictly speaking, I think the argument checking itself doesn’t allocate, so maybe there should be two keyword arguments: `check_args` and `sort_args`?

---

<div class="post-metadata">

**Author:** ![nilshg](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/nilshg/32/2283_2.png) [@nilshg](https://discourse.julialang.org/u/nilshg)\
**Post date:** [July 18, 2023, 10:19am UTC](https://discourse.julialang.org/t/is-allocation-inevitable-when-generating-random-numbers-from-a-categorical-distribution/101718/4 "2023-07-18T10:19:00Z")

</div>

Right, in that case the call to

```julia
rand!(Categorical(w, check_args = false), tmpI)

```

would have the same runtime and 8 allocations like

```julia
rand!(cat_dist, tmpI)

```

above, but the 8 allocations from calling `rand!` one a categorical distribution remain.

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [July 18, 2023, 10:25am UTC](https://discourse.julialang.org/t/is-allocation-inevitable-when-generating-random-numbers-from-a-categorical-distribution/101718/5 "2023-07-18T10:25:47Z")

</div>

Oh right, I had not paid enough attention, and I went straight for the problem I had already encountered. Time to dig some more

---

<div class="post-metadata">

**Author:** ![algunion](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/algunion/32/51630_2.png) [@algunion](https://discourse.julialang.org/u/algunion)\
**Post date:** [July 18, 2023, 10:31am UTC](https://discourse.julialang.org/t/is-allocation-inevitable-when-generating-random-numbers-from-a-categorical-distribution/101718/6 "2023-07-18T10:31:13Z")

</div>

Hope I am not operating on wrong assumptions here, but there are no allocations when calling `rand(cat_dist)`. At first glance, and maybe naive, I would then expect the `rand!` to perform similarly as the other distributions (e.g., no allocations for the in-place `rand!`):

```julia
function alloctest()
    w = rand(10^5)
    Distributions.normalize!(w,1)    
    c = Categorical(w, check_args=false)
    @btime rand($c)
end

alloctest()
# 98.000 ns (0 allocations: 0 bytes)

```

It is obvious that the output of `rand` will 0-allocate. But the `rand!` call seems to perform additional allocations before producing the output for `rand` call.

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [July 18, 2023, 10:36am UTC](https://discourse.julialang.org/t/is-allocation-inevitable-when-generating-random-numbers-from-a-categorical-distribution/101718/7 "2023-07-18T10:36:08Z")

</div>

Ok, so I did some more digging. To sample from `d isa Categorical`, the package first creates `s = sampler(d)` which is an `AliasTable`, and then calls `rand(rng, s)` as many times as necessary.

Lo and behold the `AliasTable` [implementation](https://github.com/JuliaStats/Distributions.jl/blob/master/src/samplers/aliastable.jl):

```julia
struct AliasTable <: Sampleable{Univariate,Discrete}
    accept::Vector{Float64}
    alias::Vector{Int}
end
ncategories(s::AliasTable) = length(s.alias)

function AliasTable(probs::AbstractVector)
    n = length(probs)
    n > 0 || throw(ArgumentError("The input probability vector is empty."))
    accp = Vector{Float64}(undef, n)
    alias = Vector{Int}(undef, n)
    StatsBase.make_alias_table!(probs, 1.0, accp, alias)
    AliasTable(accp, alias)
end

function rand(rng::AbstractRNG, s::AliasTable)
    i = rand(rng, 1:length(s.alias)) % Int
    u = rand(rng)
    @inbounds r = u < s.accept[i] ? i : s.alias[i]
    r
end

```

So the sampling itself is allocation-free, but the construction is not. I think what you need to do is amortize it by constructing `s = sampler(d)` at the beginning of a loop and then use the sampler directly as `x[i] = rand(rng, s)` for each component of the target `x`.

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [July 18, 2023, 10:42am UTC](https://discourse.julialang.org/t/is-allocation-inevitable-when-generating-random-numbers-from-a-categorical-distribution/101718/8 "2023-07-18T10:42:12Z")

</div>

The solution:

```julia
julia> using Distributions, LinearAlgebra

julia> function loop()
           w = rand(10^5)
           normalize!(w,1)
           tmpI = zeros(Int,10^5)
           d = Categorical(w, check_args=false) # this doesn't allocate
           s = sampler(d) # this allocates
           @time for i in 1:10^3
               for j in eachindex(tmpI) # can't use rand! directly, sad emoji
                   tmpI[j] = rand(s) # this doesn't allocate
               end
           end
       end
loop (generic function with 1 method)

julia> loop()
  2.891325 seconds

```

---

<div class="post-metadata">

**Author:** ![taka255](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/taka255/32/51426_2.png) [@taka255](https://discourse.julialang.org/u/taka255)\
**Post date:** [July 18, 2023, 10:54am UTC](https://discourse.julialang.org/t/is-allocation-inevitable-when-generating-random-numbers-from-a-categorical-distribution/101718/9 "2023-07-18T10:54:07Z")

</div>

Thank you all for your help! It was a big trap for me.

---

<div class="post-metadata">

**Author:** ![algunion](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/algunion/32/51630_2.png) [@algunion](https://discourse.julialang.org/u/algunion)\
**Post date:** [July 18, 2023, 12:14pm UTC](https://discourse.julialang.org/t/is-allocation-inevitable-when-generating-random-numbers-from-a-categorical-distribution/101718/10 "2023-07-18T12:14:54Z")

</div>

> [@gdalle](#):
>
> ```julia
> @time for i in 1:10^3
> for j in eachindex(tmpI) # can't use rand! directly, sad emoji
> tmpI[j] = rand(s) # this doesn't allocate
> end
> end
> 
> ```

@gdalle, @taka255,

I propose this version which actually supports the usage of `rand!` directly:

```julia
using Random, Distributions, LinearAlgebra

function loop()
    w = rand(10^5)
    normalize!(w, 1)
    tmpI = zeros(Int, 10^5)
    d = Categorical(w, check_args=false) # this doesn't allocate
    s = sampler(d)
    @time for i in 1:10^3
        rand!(s, tmpI)
    end
end

loop()
# 1.841305 seconds
# this is not faster than the version proposed by @gdalle
# I also get faster times for his version on this machine

```

This works because (interpret this in the context/scope of the `loop` function):

```julia
s = sampler(d)
@info d isa Sampleable # true
@info s isa Sampleable # true

#and
sampler(d) # allocates
sampler(s) # does not allocate 

#and
@info s === sampler(s) # true
# just to add redundant stuff:
@info sampler(s) === identity(s) # true

```

And there is a fallback method with the signature: `rand!(rng::AbstractRNG, s::Sampleable{Univariate}, A::AbstractArray{<:Real})` - but hey, `sample` returns a `Sampleable`. You can go into more details [here](https://github.com/JuliaStats/Distributions.jl/blob/c8d3e4b52ea2c04f655510af85f93fa3876f258f/src/univariates.jl#L140C5-L140C5).

So, the `rand!(::Categorical ,...)` method used initially by OP has a fallback on the above `rand!` method I used in the updated `loop` - which is also exported by the API - so safe to use.

Anyway - I hope my explanation is clear enough - if it is not, please run the `loop` function from above first, and we can talk later 🙂