Avoiding type instability when slicing a tuple

4
votes

This question is an adapted version of what appeared here in the JuliaLang Zulip helpdesk.

Suppose I have a function that takes in a heterogeneously typed Tuple and will return a slice of that tuple where the slice indices may be statically inferred from only type information. How can I write my function in such a way that the output type is correctly inferred?

For example, suppose my function is

function f(t::Tuple, A::Array{T, N}) where {T, N}
    if T <: AbstractFloat
        imin = 1
    elseif T <: Integer
        imin = 2
    else
        imin = 3
    end
    imax = N+2    
    t[imin:imax]
end

we see that type inference only figures out that this produces a Tuple, not it's length or element types even though all the needed information is available at compile time?

julia> let t = (:a, "b", 2, 3.0, Val(1), 2+im), A = rand(Int, 3,3)
           Base.return_types(f, Tuple{typeof(t), typeof(A)})
       end
1-element Array{Any,1}:
 Tuple

How can I write f such that this works?

2
julia

2 Answers

5
votes

The strategy I'd be most comfortable with (but maybe there's an easier way?) is to write a @generated function to manually ensure julia does the type level operations I want at compile time:

@generated function f2(t::Tuple, A::Array{T, N}) where {T, N}
    if T <: AbstractFloat
        imin = 1
    elseif T <: Integer
        imin = 2
    else
        imin = 3
    end
    imax = N+2
    out_expr = Expr(:tuple, (:(t[$i]) for i ∈ imin:imax)...)
end

The idea here is that in the generated function body, at compile time, we determine what imin and imax are, and then we manually build out an expression for our function body that reads (t[imin], t[imin+1], ..., t[imax-1], t[imax]).

For whatever reasons, julia is better able to reason about a sequence of getindex(::Tuple, ::Int) than it is about slicing a tuple, even with a statically known slice, so by manually building this expression, the compiler is able to do what we want:

julia> let t = (:a, "b", 2, 3.0, Val(1), 2+im), A = rand(Int, 3,3)
           Base.return_types(f2, Tuple{typeof(t), typeof(A)})
       end
1-element Array{Any,1}:
 Tuple{String,Int64,Float64}

Voila, the inferred output type is aTuple of length 3 whose elements are statically known to be a String, and Int and Float64!

1
votes

You have to call some unexported functions, but you can do it without a @generated function like this:

julia> function f(t::Tuple, A::Array{T, N}) where {T, N}
           t = Base.IteratorsMD.split(t, Val(N+2))[1]
           if !(T<:AbstractFloat)
               t = Base.tail(t)
               if !(T<:Integer)
                   t = Base.tail(t)
               end
           end
           return t
       end
f (generic function with 1 method)

julia> let t = (:a, "b", 2, 3.0, Val(1), 2+im), A = rand(Int, 3,3)
           Base.return_types(f, Tuple{typeof(t), typeof(A)})
       end
1-element Array{Any,1}:
 Tuple{String,Int64,Float64}

There are advantages in doing this with inferrable tuple manipulations: your code will likely compile faster and it's also Revise-able.