Type instability of discrete mixture distribution

General Usage
,
1

Hi all,

I want to define a custom distribution that is a modification of the Geometric distribution (the one that has support {1,2,3,...}) with an additional “atom” at the event 0. I.e. this is a well-defined distribution on the natural numbers {0, 1, 2, ...}. The code below correctly constructs this distribution in Julia using a mixture model from Distributions.jl, which is already quite nice.

However, I noticed a type instability when sampling random numbers from that distribution. It seems this currently slows down my code, and I don’t understand what’s going on.

using Distributions

a = 0.2
b = 0.995

d1 = Dirac(0)
d2 = Geometric(1-b) + 1
d = MixtureModel([d1, d2], [a, 1-a])

@code_warntype rand(d)

The @code_warntype output reads:

MethodInstance for rand(::MixtureModel{Univariate, Discrete, Distribution{Univariate, Discrete}, Categorical{Float64, Vector{Float64}}})
  from rand(s::Sampleable, dims::Int64...) @ Distributions ~/.julia/packages/Distributions/UaWBm/src/genericrand.jl:22
Arguments
  #self#::Core.Const(rand)
  s::MixtureModel{Univariate, Discrete, Distribution{Univariate, Discrete}, Categorical{Float64, Vector{Float64}}}
  dims::Tuple{}
Body::Any
1 ─ %1 = Distributions.rand::Core.Const(rand)
│   %2 = Distributions.default_rng()::Core.Const(Random.TaskLocalRNG())
│   %3 = Core.tuple(%2, s)::Tuple{Random.TaskLocalRNG, MixtureModel{Univariate, Discrete, Distribution{Univariate, Discrete}, Categorical{Float64, Vector{Float64}}}}
│   %4 = Core._apply_iterate(Base.iterate, %1, %3, dims)::Any
└──      return %4

Any help really appreciated!

2

The reason is that you’re building a mixture from two different distribution types, namely Dirac and Geometric.
Do you need the mixture for something else than sampling?
If you only need it for sampling, the easiest solution would be to compute the mixture manually: always sample a geometric random variable, and then with some probability return 0 instead.

3

Indeed, using a mixture of [d1 d1] or [d2 d2] produces type stable code. But I do not really understand why different distributions means type instable code, it really does not need to be and maybe should be reported as a bug ?

Edit: the fact that we can draft a type stable version clearly means that this can be considered a bug.

4

Hmm okay, I don’t see a reason why it should not work in principle. So it’s more like that this feature is not implemented?

Yes sampling this manually is of course possible, although I would sample the Dirac branch first. Then sometimes one can skip the Geometric sample.

5

Looks like rand(d.components) is type unstable, and this is logic. But then:

julia> @code_warntype rand.(d.components)
MethodInstance for (::var"##dotfunction#225#1")(::Vector{Distribution{Univariate, Discrete}})
  from (::var"##dotfunction#225#1")(x1) @ Main none:0
Arguments
  #self#::Core.Const(var"##dotfunction#225#1"())
  x1::Vector{Distribution{Univariate, Discrete}}
Body::AbstractVector
1 ─ %1 = Base.broadcasted(Main.rand, x1)::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, typeof(rand), Tuple{Vector{Distribution{Univariate, Discrete}}}}
│   %2 = Base.materialize(%1)::AbstractVector
└──      return %2

but

julia> typeof(rand.(d.components))
Vector{Int64} (alias for Array{Int64, 1})

julia>

which is weird. I’m sure this can be fixed in Distributions.jl’s code.

6

Yes, multiple samples at once, e.g. rand(d, 10) seem type stable.

So this feels more like a bug? If yes I’m happy to report this as issue.

7

Yep i think this is the right thing to do, for me its a bug / at least a missing feature.

1 Like
8

I reported this issue on GitHub. The bug (or missing feature) is actually quite general, also applying to mixtures of continuous distributions (of different types).

9

Distribution.jl defines a separate datatype for each distribution type. As a consequence, the syntax [d1 d2] (put this two varaibles of different types into an array) is resulting in an array of Any. Then to the question at compile time “what is the type of a[i]” where a is the above array, the compiler says “could be Anything”. And voilà, type instability.

1 Like
10

No, I think this is simply not the case.

[d1, d2] is of type Vector{Distribution{Univariate, Discrete}}, and also @code_warntype [d1, d2] does not report issues to me.

EDIT: I checked some more cases:
Indeed @code_warntype [d1, d2][2] does report a type instability.
However @code_warntype rand([d1, d2][2]) is type stable again. This is super confusing to me :smiley:

11

Hi,

If you look at the code Distributions.jl/src/univariate/discrete/geometric.jl at master · JuliaStats/Distributions.jl · GitHub
line 24: Geometric is a concrete subtype of the Abstract type DiscreteUnivariateDistribution.

So even though [d1 d2] is an array of Distribution{Univariate, Discrete}, (OK, not of Any :slight_smile: ), this is still not a concrete type.

For example, Vector{Real} will cause type instabilities, because on accessing, you may get a Float64, a Float32 or more exotic stuff (all of which are concrete subtypes of Real, and the compiler must hedge its bets.

With a Vector{Float64}, where Float64 is a concrete type, the compiler knows what it’s going to get when accessing the array: Float64, a specific way of coding the data into bits. No messing around with ifs, you get fast code.

EDIT
By the way, there’s a blog on the subject of type stability that you might find helpful (shameless self promotion, I know :smiley: but still, I believe this will help you.)
https://www.juliabloggers.com/writing-type-stable-julia-code/

1 Like
12

ok thank you! This explains some of the behaviour I guess.

So it may be very difficult to infer a concrete type of such arrays of different distributions. Although I still think theoretically this should be possible when the distributions have the same value support…

If there is no strong objection, I would maybe leave the GitHub issue there and see what they say…?

13

I would say your issue is not a bug, but a feature: it’s a consequence of Distributions.jl introducing different types for different distributions and providing an API encouraging you to use arrays of various distributions.

I have not used this package, but if they had a syntax like
MixtureModel(d1,d2,[a,1-a]), that could be made typestable:
You give it two distributions, the compiler says “OK, compile a method instance of MixtureModel with first parameter of type Dirac{Float64} and second of type Geometric{Float64}, concrete stuff that.”

Understanding the type system was for med the critical point of “making Julia work”. You are not alone! :smiley:

1 Like
14

Ok thanks, maybe there is a way to make it possible. I think it would be a very nice feature to have.

I still don’t get why @code_warntype rand(d) (type instable, Body::Any) and @code_warntype rand(d, 1) (type stable, Body::Vector{Int64}) behave like this?

15

These are probably methods of rand defined by Distributions.jl, a package I haven’t used - yet. Someone else will have to help.

16

Maybe the underlying storage of Mixtures should be tuples and not arrays then, to allow for type stability. But then for large mixtures there would be performance issues