How to annotate function to return 1D array of the same type as the input 2D array?

How do I express something like this:

julia> function f(mat₁::T{2}, mat₂::T{2})::T{1} where {T <: DenseArray{<:Real,N} where N}
          return eachcol(mat₁) .⋅ eachcol(mat₂)
       end
ERROR: TypeError: in Type{...} expression, expected UnionAll, got TypeVar

Basically, I want to say that function takes two matrices both of type T, which is a subtype of DenseArray{<:Real,2} and returns a vector of the same array type i.e T{1}?

For example if the input is CuMatrixes then return CuVector.

There is no generic way to go from a type which is <: Array{T, 2} to the “same” or equivalent type which is <: Array{T, 1}. After all, for some arbitrary matrix-like type M, there is no reason to expect that there is such an equivalent type (or if there is, that it is unique).

The easiest thing to do here would be to drop the return type annotation. They’re rarely necessary in Julia, and most of the time you can remove them with no effect on performance.

If that’s not an option for some reason, then perhaps you could do something like:

julia> vector_like(::Type{Array{T, 2}}) where {T} = Array{T, 1}
vector_like (generic function with 2 methods)

julia> vector_like(T) = Any # fallback for everything else
vector_like (generic function with 2 methods)

julia> function foo(x::T)::vector_like(T) where {T}
         x[:, 1]
       end
foo (generic function with 1 method)

julia> foo(zeros(2, 2))
2-element Array{Float64,1}:
 0.0
 0.0

julia> using SparseArrays

julia> foo(spzeros(2, 2))
2-element SparseVector{Float64,Int64} with 0 stored entries

where the vector_like(::Type{CuMatrix{T}}) could also return CuVector{T}. But again, it’s not obvious that this will actually be better than just omitting the return type entirely.

2 Likes