Reducing dynamic dispatch using type inference?

First of all, apologies for the terrible title, but I’m having a hard time describing what I want to do.

I’ve written a lot of my code for an optimisation problem to look something like this:

using JuMP

# Function definitions
function get_variable_x(m::Model)
    x = get(m.ext, :x, nothing)
    return get_variable_x(m, x)
end

function get_variable_x(m::Model, x::Nothing)
    return m.ext[:x] = @variable(m, x[i=1:10] >= 0)
end

get_variable_x(m::Model, x) = x

In essence, most of my functions first check if a Julia variable (in this case a JuMP optimisation variable) has been created before, returns it if so and creates it if not. The problem is that the function get_variable_x is type unstable (I think I’m using that word correctly):

using Test
m = Model()

@inferred get_variable_x(m)

This produces the following error:

return type Array{VariableRef,1} does not match inferred return type Any
error(::String) at error.jl:33
top-level scope at test.jl:19

If I simply created the variable, this wouldn’t be a problem:

function get_variable_x_alt(m::Model)
    return m.ext[:x] = @variable(m, x[i=1:10] >= 0)
end

m = Model()

@inferred get_variable_x_alt(m) # no error

Of course, I could further constrain the types that get_variable_x can by writing this:

get_variable_x(m::Model, x::Array{VariableRef,1}) = x

However, I would like the Julia’s type inference to do the job for me. I’m sure that there is an elegant way of doing this, but I just can’t think of it.

It seems to me you can do this just as well without involving dispatch

function get_variable_x(m::Model)
    # could also use get-do syntax; whatever you fancy
    if !haskey(object_dictionary(m), :x)
        @variable(m, @variable(m, x[i=1:10] >= 0)
    else
        x = m[:x]
    end
    return x
end

I think this procedure is inherently not type stable though, because the key :x can map to anything (the dict is Symbol=>Any). This will be true regardless of how exactly you do it by virtue of how the dictionaries in the Model are defined. You could define your own dictionary or struct that maps Symbol => whatever to resolve this, or if you know exactly what each of your symbols will represent, you could constrain the return type of each “getter” function so at least at the calling point the return type can be known (note that internally the getter is still unstable). E.g. in the definition:

function get_variable_x(m::Model)::Vector{VariableRef}
    ...

I think it’s definitely worth asking whether this is a worthwhile optimization to make. Does this type instability matter for what you’re doing?

1 Like

I was actually doing it similarly to how you wrote it before until I read some part of the Julia documentation about multiple dispatch.

Honestly I don’t think this is slowing down my code significantly - it seems to run pretty fast already (at least with respect to solving the optimisation problem in Model). I’m mainly asking this out of interest and to satisfy my OCD.

So is there no way of implicitly specifying / inferring the return type of get_variable_x in my case? I should mention that I don’t want to specify it explicitly since they can generally be more complicated, e.g.

JuMP.Containers.DenseAxisArray{JuMP.GenericAffExpr{Float64,JuMP.VariableRef},2,Tuple{Array{Tuple{String,String},1},Array{Int64,1}},Tuple{Dict{Tuple{String,String},Int64},Dict{Int64,Int64}}}

If I changed anything in my function I would also change the type and quickly run into a lot of problems. Again though, I’m asking purely to satisfy my curiosity.

So is there no way of implicitly specifying / inferring the return type of get_variable_x in my case?

It seems not. Or at least that doing so would be extremely cumbersome, and likely slower than allowing the instability. As long as you use a function barrier between getting x and using it, you should be ok. I.e. once you have x, you should call some_func(x) which is type stable. The idea is explained well in this blog post.

1 Like

Thanks for that, I think I finally understand function barriers. I use x in JuMP optimisation constraints, so I guess that counts? i.e. when I call this:

@constraint(m, [i=1:10], i*x[i] >= 1)

This isn’t a problem then since @constraint (or at least the functions it then calls) know the type of x.