Multiple dispatch on StructArrays

Hi everyone. I am working with the StructArrays.jl package and trying to program a function to react differently for different types of arrays inside of a StructArray.

Specifically I want to differentiate between Array and CuArray of ComplexF32 values. I have found the most straighforward (and honestly ugly looking) way to do this.

# Creating the arrays
cpuarray = Array{ComplexF32}(undef, 2)
gpuarray = CuArray(cpuarray)
scpuarray = StructArray{ComplexF32}((real(cpuarray),imag(cpuarray)))
sgpuarray = StructArray{ComplexF32}((real(gpuarray),imag(gpuarray)))

# Checking types
julia> typeof(scpuarray)
StructArray{Complex{Float32},1,NamedTuple{(:re, :im),Tuple{Array{Float32,1},Array{Float32,1}}},Int64}
julia> typeof(sgpuarray)
StructArray{Complex{Float32},1,NamedTuple{(:re, :im),Tuple{CuArray{Float32,1,Nothing},CuArray{Float32,1,Nothing}}},Int64}

# Defining a basic function for both scpuarray and sgpuarray 
function example(
    input::StructArray{Complex{Float32},1,NamedTuple{(:re, :im),Tuple{Array{Float32,1},Array{Float32,1}}},Int64}
)
    return 1
end
function example(
    input::StructArray{Complex{Float32},1,NamedTuple{(:re, :im),Tuple{CuArray{Float32,1,Nothing},CuArray{Float32,1,Nothing}}},Int64}
)
    return 2
end

# Checking the definitions
julia> example(scpuarray)
1
julia> example(sgpuarray)
2

Are there any other ideas on how achieve this in a more elegant way?
Ideally I would just like to write something like input::StructArray(CuArray)
Would appreciate any help.

Are you concerned only with the length of the type names? You could write some type aliases if there aren’t too many of them.

3 Likes

Dispatch on a Trait?

1 Like

Sort of. I feel like the type name could have been more “readable” if you wish. I guess that’s my issue with it. Just like I mentioned before, it would have been nice to see something along the lines of:

# Pseudocode
function example(
       input::StructArray{ComplexF32}(Array, Array)
)
       return 1
end

You mean like this?

function example(
           input::StructArray{Complex{Float32},1,A,Int64}
       ) where {
       A <: NamedTuple{(:re, :im),Tuple{Array{Float32,1},Array{Float32,1}}}
       }
           return 1
       end
julia> function example(
           input::StructArray{Complex{Float32},1,A,Int64}
       ) where {
       A <: NamedTuple{(:re, :im),Tuple{Array{Float32,1},Array{Float32,1}}}
       }
           return 1
       end

Not really familiar with this concept. Is this what you’re mentioning? https://github.com/mauro3/SimpleTraits.jl
Would be very kind of you to link some resources, I will look into it.

Something like that. Or you could define the alias somewhere else in the code:

B = NamedTuple{(:re, :im),Tuple{Array{Float32,1},Array{Float32,1}}}
function example2(input::StructArray{Complex{Float32},1,B,Int64}) where B
    return 1
end

or define a type alias for the whole thing:

StructCpuArray  = StructArray{Complex{Float32},1,NamedTuple{(:re, :im),Tuple{Array{Float32,1},Array{Float32,1}}},Int64}
function example3(input::StructCpuArray)
    return 1
end
1 Like

Thanks! The second option is probably the most obvious and pretty one.

1 Like

Yes, SimpleTraits.jl is a good package for traits. BinaryTraits.jl is also very easy to use.

The drawback of the type alias is that you essential dispatch on implementation details of StructArrays.jl (the NamedTuple etc). That looks brittle to me.

2 Likes