Type instability in random data with specified distribution

user_8636268 0
Viewed 49

I want to generate data from a linear model with noise (Y = X*w + e) where I can specify the distributions of input vector X and scalar noise e. For this, I specify the below struct

using Distributions

struct NoisyLinearDataGenerator
    x_dist::ContinuousMultivariateDistribution
    noise_dist::ContinuousUnivariateDistribution
    weights::Vector{Float64}
end

and a function to generate N points from it:

function generate(nl::NoisyLinearDataGenerator, N)
    x = rand(nl.x_dist, N)'
    e = rand(nl.noise_dist, N)
    return x, x*nl.weights + e
end

This seems to be working, however not type stable, as

nl = NoisyLinearDataGenerator(MvNormal(5, 1.0), Normal(), ones(5))

@code_warntype generate(nl,1)

yields

Variables
  #self#::Core.Compiler.Const(generate, false)
  nl::NoisyLinearDataGenerator
  N::Int64
  x::Any
  e::Any

Body::Tuple{Any,Any}
1 ─ %1  = Base.getproperty(nl, :x_dist)::Distribution{Multivariate,Continuous}
│   %2  = Main.rand(%1, N)::Any
│         (x = Base.adjoint(%2))
│   %4  = Base.getproperty(nl, :noise_dist)::Distribution{Univariate,Continuous}
│         (e = Main.rand(%4, N))
│   %6  = x::Any
│   %7  = x::Any
│   %8  = Base.getproperty(nl, :weights)::Array{Float64,1}
│   %9  = (%7 * %8)::Any
│   %10 = (%9 + e)::Any
│   %11 = Core.tuple(%6, %10)::Tuple{Any,Any}
└──       return %11

I am not sure why this is, since I would expect the type of the sampled data to be specified by using ContinuousMultivariateDistribution and ContinuousUnivariateDistribution.

What is leading to type instability here and what should a type stable implementation look like?

julia
1 Answers

The problem is that ContinuousMultivariateDistribution and ContinuousUnivariateDistribution are abstract types. While your knowledge of statistics tells you that they probably should return Float64, there is no guarantee on a language level that someone won't implement, say, a ContinuousUnivariateDistribution that returns some other object. Therefore the compiler can't know that all ContinuousUnivariateDistribution produces any particular type.

For example, I might write:

struct BadDistribution <: ContinuousUnivariateDistribution end
Base.rand(::BadDistribution, ::Integer) = nothing

Now, you could make a NoisyLinearDataGenerator containing a BadDistribution as x_dist. What would the output type be then?

In other words, the output of generate simply can't be predicted only from its input types.

To solve this, you need to either specify specific distributions for your new type, or else make your new type parametric. In Julia, whenever we have a field of a type that cannot be specified to a concrete type, we usually leave it as a type parameter. Thus, one possible solution is this:

using Distributions

struct NoisyLinearDataGenerator{X,N}
    x_dist::X
    noise_dist::N
    weights::Vector{Float64}

    function NoisyLinearDataGenerator{X,N}(x::X, n::N, w::Vector{Float64}) where {
                                    X <: ContinuousMultivariateDistribution,
                                    N <: ContinuousUnivariateDistribution}
        return new{X,N}(x,n,w)
    end
end

function NoisyLinearDataGenerator(x::X, n::N, w::Vector{Float64}) where {
                                X <: ContinuousMultivariateDistribution,
                                N <: ContinuousUnivariateDistribution}
    return NoisyLinearDataGenerator{X,N}(x,n,w)
end

function generate(nl::NoisyLinearDataGenerator, N)
    x = rand(nl.x_dist, N)'
    e = rand(nl.noise_dist, N)
    return x, x*nl.weights + e
end
nl = NoisyLinearDataGenerator(MvNormal(5, 1.0), Normal(), ones(5))

Here, the type of nl is NoisyLinearDataGenerator{MvNormal{Float64,PDMats.ScalMat{Float64},FillArrays.Zeros{Float64,1,Tuple{Base.OneTo{Int64}}}},Normal{Float64}} (yes, I know, awful to read), but its type contain all information needed for the compiler to fully predict the output type of generate.

user_109926670
Related