Type-Instability Caused by Wrapping Type-Stable Function

The following results are puzzling to me. I have a type-stable function eval_logprior stored in structure prior_struct which I type as ::Function.

Input:

eval_logprior = prior_struct.eval_logprior
@code_warntype eval_logprior(θ)

Output:

Variables
  #self#::SimMdlPrices.var"#eval_logprior#48"{FullNormal, FullNormal, Vector{Exponential{Float64}}, Vector{Int64}, Vector{Int64}, Vector{Int64}, Vector{Int64}, Vector{Function}, SimMdlPrices.var"#66#87"{SimMdlPrices.var"#65#86"{SimMdlPrices.var"#h_fn#85"{Vector{Function}}, SimMdlPrices.var"#h_fn#85"{Vector{Function}}}}}
  θ::Matrix{Float64}
  logprior::Float64
  logp_exp::Float64
  logp_n::Float64
  logp_mvn::Float64

Body::Float64
1 ─ %1  = Core.getfield(#self#, :d_mvn)::FullNormal
│   %2  = Core.getfield(#self#, :f_array)::Vector{Function}
│   %3  = Core.getfield(#self#, :mvn_cap_ndx)::Vector{Int64}
│   %4  = Core.getfield(#self#, :mvn_tsm_ndx)::Vector{Int64}
│         (logp_mvn = SimMdlPrices.logpdf_mvn(θ, %1, %2, %3, %4))
│   %6  = Core.getfield(#self#, :d_n)::FullNormal
│   %7  = Core.getfield(#self#, :f_array)::Vector{Function}
│   %8  = Core.getfield(#self#, :n_ndx)::Vector{Int64}
│         (logp_n = SimMdlPrices.logpdf_n(θ, %6, %7, %8))
│   %10 = Core.getfield(#self#, :d_exp)::Vector{Exponential{Float64}}
│   %11 = Core.getfield(#self#, :f_array)::Vector{Function}
│   %12 = Core.getfield(#self#, :exp_ndx)::Vector{Int64}
│         (logp_exp = SimMdlPrices.logpdf_exp(θ, %10, %11, %12))
│   %14 = logp_mvn::Float64
│   %15 = logp_n::Float64
│   %16 = logp_exp::Float64
│   %17 = Core.getfield(#self#, :absdetjac_fn)::SimMdlPrices.var"#66#87"{SimMdlPrices.var"#65#86"{SimMdlPrices.var"#h_fn#85"{Vector{Function}}, SimMdlPrices.var"#h_fn#85"{Vector{Function}}}}
│   %18 = (%17)(θ)::Float64
│   %19 = SimMdlPrices.log(%18)::Float64
│         (logprior = %14 + %15 + %16 + %19)
└──       return logprior

The act of simply wrapping the function as a closure in another (practically for optimization over certain dimensions), results in a type-instability.

Input:

f(x) = prior_struct.eval_logprior(x)
@code_warntype f(θ)

Output:


Variables
  #self#::Core.Const(f)
  x::Matrix{Float64}

Body::Any
1 ─ %1 = Base.getproperty(Main.prior_struct, :eval_logprior)::Any
│   %2 = (%1)(x)::Any
└──      return %2

Any words of advice?

This line shows that, within f, it is unknown what prior_struct.eval_logprior is. This is (almost certainly, although your example was incomplete so some chance I’m wrong) because prior_struct is a non-constant global variable.

If you wrote this same code inside a function (or some other place where prior_struct was of known type), it would probably be fine.

What might work in your above example (in the REPL) is to instead define f as

f = let eval_logprior=prior_struct.eval_logprior
    (x)->eval_logprior(x)
end

which is essentially what worked for you the first time. Although getting this into a local scope is the better solution. I didn’t get to verify this so I might have made a minor mistake.

BEGIN EDIT:

The other issue that might might be in play is that you might have written

struct Prior_Struct_Type
    # stuff
    eval_logprior::Function
    # stuff
end

Function is /not/ a concrete type (it is an abstract type) so this field is non-concretely typed and any access will be type-unstable. Instead, you can use a parametric type

struct Prior_Struct_Type{F}
# or Prior_Struct_Type{F<:Function} to require that eval_logprior is a Function, but this has no performance benefit
    # stuff
    eval_logprior::F
    # stuff
end

which will be stable.

5 Likes

Thanks! A combination of the two worked perfectly :slight_smile: Much appreciated.