Obtaining the ND version of a 1D type

I have an AbstractVector type x and want to get the type of the 2D version. For example, if I’m passed a x::CuArray{T,1,device}, I’d like to end up with a type CuArray{T,2,device}.

One (type-unstable?) way to do this is to get the type of a reshape.

_promote_1D_vector_type(x::AbstractVector) = typeof(reshape(x, (length(x), 1)))

Is there a type-stable way to do this?

Do you need it to work for all abstract vectors, or would cuda vectors and julia vectors be enough? If so, you could just dispatch to those two types of vectors

_promote_1D_vector_type(x::CuArray{T,1,M}) where {T,M} = CuArray{T,2,M}
_promote_1D_vector_type(x::Vector{T}) where T = Matrix{T}

Otherwise, I think what you have is going to be best combination of simple and general. I would maybe write it as

_promote_1D_vector_type(x::AbstractVector) = typeof(reshape(x, :, 1))

because I think reshape has some specializations for colon indexing (I might be wrong here).

Also when I run @code_warntype I don’t see any type instability. Was there something else you were using that reported type unstable code?

2 Likes

Note that the reshaping approach has limitations, because it will often wrap the vector inside a ReshapedArray instead of returning the “true” 2D type, even when that type exists.
Here’s an example:

julia> using SparseArrays

julia> x = spzeros(4)
4-element SparseVector{Float64, Int64} with 0 stored entries

julia> reshape(x, 2, 2)
2×2 reshape(::SparseVector{Float64, Int64}, 2, 2) with eltype Float64:
 0.0  0.0
 0.0  0.0

julia> typeof(ans)  # not a SparseMatrixCSC
Base.ReshapedArray{Float64, 2, SparseVector{Float64, Int64}, Tuple{}}
4 Likes

I think @gdalle’s example illustrates another layer of conplexity hiding in this seemingly simple question: There just might not be a “N+1 D” version of a “N D” type.
So the best course of action is probably to use reshape as a default implementation and add methods for specific types where this does not give a good result (like SparseVector, SparseMatrixCSC,…)

4 Likes

Note that this is perfectly fine and type stable:

julia> @code_warntype _promote_1D_vector_type([1,2,3])
MethodInstance for _promote_1D_vector_type(::Vector{Int64})
  from _promote_1D_vector_type(x::AbstractVector) @ Main REPL[1]:1
Arguments
  #self#::Core.Const(Main._promote_1D_vector_type)
  x::Vector{Int64}
Body::Type{Matrix{Int64}}
1 ─ %1 = Main.typeof::Core.Const(typeof)
│   %2 = Main.reshape::Core.Const(reshape)
│   %3 = Main.length::Core.Const(length)
│   %4 = (%3)(x)::Int64
│   %5 = Core.tuple(%4, 1)::Core.PartialStruct(Tuple{Int64, Int64}, Any[Int64, Core.Const(1)])
│   %6 = (%2)(x, %5)::Matrix{Int64}
│   %7 = (%1)(%6)::Core.Const(Matrix{Int64})
└──      return %7

In NDTools.jl we have something called reorient which works for N-d arrays as it uses a Val(d) parameter. However, this is only type stable if the d comes from type e.g. the dimensionality of another array:

1 Like

Typically you use similar to construct an instance of a (mutable) similar type, e.g. similar(x, 3, 4) makes a 3 \times 4 uninitialized array of a similar type to x, which you can then initialize to desired values.

Is there a reason why you need to construct the type by itself, rather than directly constructing an instance of the type?

2 Likes

Thank you everyone for these insightful comments, I didn’t appreciate these nuances!

I’m pleasantly surprised to be wrong about the type stability, I had assumed that the reshape would cause inference since it is operating on the instance.

Here’s the context: RandomizedPreconditioners.jl takes in a PSD system A (anything that defines mul!) and builds a Nyström sketch by evaluating A on some random vectors. For example, those random vectors could be created together in a CuArray{T,2,device}, leading to a sketch which lives on GPU. The question is how to allow the user to specify where they would like these original random vectors (and thus the preconditioner) to live.

We currently allow the user to specify a type with keyword argument S = Array{Float64,2}, which seems like a reasonable choice. However, in a related package we construct a sketch from A and only have access to a type V <: AbstractVector{T} which could be a CuArray or other. So we need to convert that vector type to the appropriate matrix type. Maybe we should have set up this interface differently?

Thank you, I’m glad to see that this function has some precedent and is type-stable! Although the example from @gdalle is somewhat troubling.

I could imagine doing this for some important vector types, but it does feel like I’m writing C++ rather than generic code…