Specialization on vararg of types

How do I get Julia to specialize to Vararg{Type}?

The docs say:

As a heuristic, Julia avoids automatically specializing on argument type parameters in three specific cases: Type, Function, and Vararg.

However, they say you can get around this for Vararg, for arbitrary types with Vararg{Any,N} where N:

julia> h_vararg(x::Vararg{Any, N}) where {N} = x
h_vararg (generic function with 1 method)

julia> Base.specializations(@which h_vararg(1f0, 1.0))
Base.MethodSpecializations(MethodInstance for h_vararg(::Float32, ::Float64))

Then they say you can specialize on Type with Type{T} where T:

julia> g_type(t::Type{T}) where T = t
g_type (generic function with 1 method)

julia> Base.specializations(@which g_type(Float32))
Base.MethodSpecializations(MethodInstance for g_type(::Type{Float32}))

But what about if I want to combine them? Let’s see:

julia> h2_vararg(x::Vararg{Any,N}) where {N} = x
h2_vararg (generic function with 1 method)

julia> Base.specializations(@which h2_vararg(Float32, Float64))
Base.MethodSpecializations(MethodInstance for h2_vararg(::Type, ::Type))

Nope, this doesn’t specialize because we didn’t do the Type{T} where T trick. But how should we do that within a Vararg? Maybe just this?

julia> h3_vararg(x::Vararg{Type,N}) where {N} = x
h3_vararg (generic function with 1 method)

julia> Base.specializations(@which h3_vararg(Float32, Float64))
Base.MethodSpecializations(MethodInstance for h3_vararg(::Type, ::Type))

Nope… So how do we actually force specialization here?


Related:

1 Like

Found a way!

julia> @generated f_vararg(t::Vararg{Type}) = :(t)
f_vararg (generic function with 1 method)

julia> Base.specializations(@which f_vararg(Float32, Float64))
Base.MethodSpecializations(MethodInstance for f_vararg(::Type{Float32}, ::Type{Float64}))

Not sure if there’s a way that avoids generated functions, but this at least gets the job done.

But please share any alternatives as I am quite curious!

julia> f_vararg(t::Vararg{Type{<:T}}) where {T} = t
f_vararg (generic function with 1 method)

julia> Base.specializations(@which f_vararg(Float32, Float64))
Base.MethodSpecializations(MethodInstance for f_vararg(::Type{Float32}, ::Type{Float64}))
4 Likes

Brilliant!

Do you have some intuition for why this works? For instance, I would have thought you would have to declare the N number for the Vararg to specialize. But it seems here it is not needed.

No idea. Just poked around to find something that worked. But I’m guessing that having a Type that depends on a declared typevar (i.e., a where clause) somewhere in the signature is sufficient to force specialization.

1 Like

Just be aware of a kinda-weird interplay between types and parameterizations here:

julia> f_vararg(t::Vararg{Type{<:T}}) where {T} = T

# ok
julia> f_vararg(Int, Int)
Int64

# specializes and generally works, but T is not defined
julia> f_vararg(Int, String)
ERROR: UndefVarError: `T` not defined
4 Likes

I don’t know if it was left out for some good reason, but the Vararg section is actually pretty vague about what is being specialized. Unlike lone functions and types, Vararg specifies multiple types with a possibly fixed arity. From what I could tell (though you should probably verify yourself), the “nonspecializing” __(x::Int...) specializes on Int just fine, it’s the arity it gives up on. Even __(x...) specializes on the first argument’s type in 2 specializations, one where the rest of the types match, and one where they are Any (or whatever abstract type annotates x). Making a method parameter __(x::T...) where T specializes on arguments matching 1 concrete type (and does not specialize on arity); this also specializes on input types like Type{Float32} already. Specifying the arity parameter, usually N, specializes over arity and every argument’s type.

That is, except for types themselves as inputs, which I still don’t understand. If consistent with non-type inputs, specializations should show their concrete types DataType (most types), Union (unions), UnionAll (iterated unions of parametric types), not the abstract Type. You say it doesn’t specialize, and for good reason, but it’s certainly not falling back to Any for the 2nd argument onward, either. It’s more like it specializes, but weirdly. Vararg{Type,N} doesn’t change anything because there’s no method parameter for type selection.

danielwe’s solution performs the type selection pattern for unlimited arguments in an unintuitive and creative way. It can be written fully as f_vararg(t::Vararg{Type{S} where S<:T}) where {T}, so it makes sense that each selected type S is a subtype of the 1 method parameter T<:Any. This also specializes over arity despite not specifying a method parameter for it, so this is indeed the equivalent to h_vararg(x::Vararg{Any, N}) where {N} for non-type inputs.

Other equivalents for completion:

  • __(t::Vararg{Type{Int}}) specifies a type but does not specialize on arity, like __(x::Int...) for non-type inputs.
  • no equivalents to __(x::T...) where T or __(x...) because a lone method parameter for the type also specializes on arity. Weird how x::T... didn’t specialize on arity too, but maybe the syntax mattered.
  • __(t::Vararg{Type{Int}, N}) where N specializes on arity and the specified type, like __(x::Vararg{Int, N}) where N for non-type inputs.
  • __(t::Vararg{Type{T}}) where {T} specializes on arity and 1 matching type, like __(x::Vararg{T,N}) where {T,N}, which also handles input types already.

In the vast majority of cases, a method parameter must be 1 known instance at compile-time for dispatch to work, even if it is absent from the body. I suppose an exception was spotted finally. Example below of the usual case:

julia> myeltype(::Type{S} where {S<:AbstractArray{T,N} where N}) where {T} = 42
myeltype (generic function with 1 method)

julia> myeltype(Union{Vector{Int}, Matrix{Int}}) # T is Int
42

julia> myeltype(Union{Vector{Int}, Vector{Bool}}) # T is unspecific
ERROR: MethodError: no method matching myeltype(::Type{Union{Vector{Bool}, Vector{Int64}}})

For posterity since this sort of compiler stuff might be version-dependent, this was all done on v1.10.0

3 Likes

NB: Test.detect_unbound_args warns about this method because the static parameter is not always fully constrained (as shown by @aplavin):

julia> f_vararg(t::Vararg{Type{<:T}}) where {T} = t
f_vararg (generic function with 1 method)

julia> using Test

julia> detect_unbound_args(Main)
[1] f_vararg(t::Type{<:T}...) where T @ Main REPL[1]:1

julia> @code_warntype f_vararg(Int, String)
MethodInstance for f_vararg(::Type{Int64}, ::Type{String})
  from f_vararg(t::Type{<:T}...) where T @ Main REPL[1]:1
Static Parameters
  T <: Union{Int64, String}
[...]
1 Like

NB: JET.jl is buggy here, reporting a false negative:

Is it possible that this unorthodox feature is also a bug because of this? Would be a shame considering we don’t have alternatives this deep into the type system.