The Broadcasting API and small Unions

General Usage
1

After trying some near examples and the docs, I have not yet made the Broadcasting API do this. More explicit guidance is appreciated.

abstract type OtherAbstractFloat <: AbstractFloat end
primitive type Float64a <: OtherAbstractFloat 64 end
primitive type Float64b <: OtherAbstractFloat 64 end

const Float64s = Union{Float64, Float64a, Float64b}
const MaybeFloat64s = Union{Missing, Float64, Float64a, Float64b}

# Float64a, Float64b work like Float64; show with '\^a', '\^b'

Base.Float64(x::T) where {T<:OtherAbstractFloat} = reinterpret(Float64, x)
Float64a(x::Float64) = reinterpret(Float64a, x)
Float64b(x::Float64) = reinterpret(Float64b, x)

Base.show(io::IO, x::Float64a) = print(io, string(Float64(x),"ᵃ"))
Base.show(io::IO, x::Float64b) = print(io, string(Float64(x),"ᵇ"))

square(x::T) where {T<:AbstractFloat} = T(Float64(x)^2)
square(x::Missing) = missing

testvec1 = Float64s[Float64(1.0), Float64a(2.0), Float64b(3.0)]
testvec2 = MaybeFloat64s[Float64(1.0), Float64a(2.0), missing]

# this happens without using the Broadcasting API

square.(testvec1)
3-element Array{AbstractFloat,1}:
 1.0 
 4.0ᵃ
 9.0ᵇ

square.(testvec2)
3-element Array{Any,1}:
 1.0     
 4.0ᵃ    
  missing

# I want to obtain

square.(testvec1) # 3-element Array{Float64s,1}
3-element Array{Union{Float64, Float64a, Float64b},1}:
 1.0 
 4.0ᵃ
 9.0ᵇ

square.(testvec2) # 3-element Array{MaybeFloat64s,1}
3-element Array{Union{Missing, Float64, Float64a, Float64b},1}:
 1.0     
 4.0ᵃ    
  missing
2

Seems that when the number of elements in the union goes over two, inference falls back to Any.

julia> Base._return_type(square, Tuple{Union{Missing, Float64}})
Union{Missing, Float64}

julia> Base._return_type(square, Tuple{Union{Missing, Float64, Float64a}})
Any
3

That seems inconsistent with all the effort done to make working with small unions (up to four types) performant. As I understand it, most of that work is about handling vectors and arrays of elements typed as small unions. I thought that the Broadcasting API could be used to get the result I seek by using it to make this logic happen:


function Base.Broadcast.broadcast(fn, x::AbstractArray{T,N}) where 
                                                             {N, T<:Float64s}
    result = similar(x)
    @inbounds @simd for idx in eachindex(view(x, axes(x)...,))
        result[idx] = fn(x[idx])
    end
    return result
end
4

Not taking into a account fn when allocating the result doesn’t seem like it would work. What if fn was x -> convert(Int, x)?

A workaround is:

res = similar(testvec2);
res .= square.(testvec2)
5

Yes, that is true. In my application I need to make arrays of small unions of Float64-like primitive types respond to floating point math functions just as arrays of Float64s do. So I know the result types for each group of functions I “delegate”.

That workaround looks very helpful. Is there a way to use it with a specialized Broadcast Style and broadcast_similar so that a client app could write squares = square.(testvec1) and generally res = mathfunction.(testvec1)?

6

In Julia 0.6:

julia> square(x::Int)=x*x
square (generic function with 2 methods)

julia> Base._return_type(square, Tuple{Union{Int64, Float64}})
Union{Float64, Int64}

julia> arr=Union{Int64,Float64}[1,2.0]
2-element Array{Union{Float64, Int64},1}:
 1  
 2.0

julia> square.(arr)
2-element Array{Real,1}:
 1  
 4.0
7

Ref. https://github.com/JuliaLang/julia/issues/27106.

8

FWIW, I was wrong here, _return_types is only used if the return type is concrete.