Preserving Array's element-type union in loops or broadcasts to avoid dynamic dispatch

Question

I want to do elementwise operations on Array{Union{A,B,C}, M}, and I want an output of Array{Union{A,B,C}, N} without dynamic dispatch caused by failure to infer the Union{A,B,C} element type.

I pasted the examples I'm playing with below. I am aware I could use type annotations to control return types e.g. local y::eltype(x), but it doesn't stop Any inferences in the elementwise operations (as shown by @code_warntype) or dynamic dispatch (won't affect allocation numbers, which I report later).

It seems like Union splitting works up to 3 concrete types in the example, but while looping splits the element type, broadcasting splits the Vector type instead of the element type parameter. I'm asking for generalizable approaches to writing loops or elementwise methods that get around these inference limitations, even if it gets as tedious as conditional statements for each type.

begin
  function loop(x)
    local y
    for i in x
      y = i
    end
    return y
  end
        
  function each1(x)
    return x+1
  end
  
  # inputs to test
  const x3 = Union{Int, Float64, Bool}[1, 1.1, true]
  const x4 = Union{Int, Float64, Bool, ComplexF64}[1, 1.1, true, 1.1im]
end

I won't paste the whole @code_warntype printouts, just the inferred output types. Also I'm not exactly sure how to measure dynamic dispatch, but it causes extra allocations so I'm reporting the allocations in the 2nd runs of @time.

typeof(loop(x3)) # Bool
@code_warntype loop(x3) # Union{Bool, Float64, Int64}
@time loop(x3) # 0 allocations

typeof(each1.(x3)) # Vector{Real}
@code_warntype each1.(x3) # Union{Vector{Float64}, Vector{Int64}, Vector{Real}}
@time each1.(x3) # 3 allocations

typeof(loop(x4)) # ComplexF64
@code_warntype loop(x4) # Any
@time loop(x4) # 8 allocations

typeof(each1.(x4)) # Vector{Number}
@code_warntype each1.(x4) # AbstractVector{<:Number}
@time each1.(x4) # 13 allocations

Answer 1

The reason for what you see it twofold:

1. as you noticed if union is larger than of 3 types the Julia compiler gives up trying to do exact type inference and uses more crude bounds. This is a hard coded assumption to avoid excessive code generation.
2. Broadcasting performs narrowing of eltype of a collection if possible.

Let me present the second issue in a minimal example:

julia> x = Any[1]
1-element Vector{Any}:
 1

julia> identity.(x)
1-element Vector{Int64}:
 1

Also when you write:

broadcasting splits the Vector type instead of the element type parameter

a precise information is that broadcasting informs you that the result will be either Vector{Float64}, or Vector{Int64}, or Vector{Real} depending on the actual contents of the array (beause eltype narrowing is performed).

The way to avoid it is to use broadcasting assignment instead (or just a loop). However, you must be sure that the eltype of the resulting container is able to store the result of the operation. So in your case it could be e.g.:

julia> r3 = similar(x3, Union{Int, Float64})
3-element Vector{Union{Float64, Int64}}:
 0.0
 0.0
 0.0

julia> r3 .= each1.(x3)
3-element Vector{Union{Float64, Int64}}:
 2
 2.1
 2

julia> @time r3 .= each1.(x3)
  0.000034 seconds (1 allocation: 16 bytes)
3-element Vector{Union{Float64, Int64}}:
 2
 2.1
 2

(note that I have manually chosen the eltype of r3 so that it is proper - i.e. I know what type of result I can potentially expect given the definition of each1 and the type of x3)