Map over vector with union eltype makes type-inference worse

julia> struct A end

julia> struct B end

julia> struct C{AA,BB}
        a :: AA
        b :: BB
       end

julia> v = Union{C{A,B}, C{B,B}}[C(A(), B()), C(B(),B())]
2-element Vector{Union{C{A, B}, C{B, B}}}:
 C{A, B}(A(), B())
 C{B, B}(B(), B())

julia> map(identity, v)
2-element Vector{C{AA, B} where AA}:
 C{A, B}(A(), B())
 C{B, B}(B(), B())

This seems to happen because

julia> Base.@default_eltype(v)
C{AA, B} where AA

How do I ensure that the eltype of the result of the map remains a small union? My actual use case is more complicated than this, and I can’t add type-assertions or such hacks.

2 Likes

I doubt this solves your problem beyond this MWE, but this specific problem does not appear for v = C{Union{A,B}, B}[C(A(), B()), C(B(), B())]. Another (ugly) option is to preallocate your result with the desired type as in map!(identity, similar(v), v), but that’s annoying to maintain as I often use map because I want the compiler to determine the output type for me (although I rarely use it for unstable functions, so don’t usually encounter your problem).

A slightly simpler set of MWEs for this is:

julia> struct A end; struct B end;

julia> map(identity, Union{A, B}[A(), B()])
2-element Vector{Any}:
 A()
 B()

julia> map(identity, Union{Int64,Float64}[1, 2e0])
2-element Vector{Real}:
 1
 2.0

So any inhomogeneous output appears to be homogenized (by type widening, although not all the way to Any if another abstract type is suitable) by map (or broadcast). It’s not the eltype of the input that matters, but the return type of the mapped function as seen here:

julia> map(x -> x < 3 ? A() : B(), 1:4)
4-element Vector{Any}:
 A()
 A()
 B()
 B()

julia> map(Returns(3//1), Union{Int64,Float64}[1, 2e0])
2-element Vector{Rational{Int64}}:
 3//1
 3//1
1 Like

OP knows this, but related thread.