Merge pull request #80 from JuliaDiffEq/concrete

fix GPU concretizations

Author: ChrisRackauckas

Project.toml

@@ -1,7 +1,7 @@
 name = "RecursiveArrayTools"
 uuid = "731186ca-8d62-57ce-b412-fbd966d074cd"
 authors = ["Chris Rackauckas <[email protected]>"]
-version = "1.2.1"
+version = "2.0.0"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"

src/RecursiveArrayTools.jl

@@ -5,8 +5,8 @@ module RecursiveArrayTools
   using Requires, RecipesBase, StaticArrays, Statistics,
         ArrayInterface
 
-  abstract type AbstractVectorOfArray{T, N} <: AbstractArray{T, N} end
-  abstract type AbstractDiffEqArray{T, N} <: AbstractVectorOfArray{T, N} end
+  abstract type AbstractVectorOfArray{T, N, A} <: AbstractArray{T, N} end
+  abstract type AbstractDiffEqArray{T, N, A} <: AbstractVectorOfArray{T, N, A} end
 
   include("utils.jl")
   include("vector_of_array.jl")

src/init.jl

@@ -10,4 +10,13 @@ function __init__()
     RecursiveArrayTools.recursive_unitless_bottom_eltype(a::ApproxFun.Fun) = recursive_unitless_bottom_eltype(ApproxFun.coefficients(a))
     RecursiveArrayTools.recursive_bottom_eltype(a::ApproxFun.Fun) = recursive_bottom_eltype(ApproxFun.coefficients(a))
   end
+
+  @require CuArrays="3a865a2d-5b23-5a0f-bc46-62713ec82fae" begin
+    function CuArrays.CuArray(VA::AbstractVectorOfArray)
+      vecs = vec.(VA.u)
+      return CuArrays.CuArray(reshape(reduce(hcat,vecs),size(VA.u[1])...,length(VA.u)))
+    end
+
+    Base.convert(::Type{<:CuArrays.CuArray},VA::AbstractVectorOfArray) = CuArrays.CuArray(VA)
+  end
 end

src/vector_of_array.jl

@@ -1,13 +1,19 @@
 # Based on code from M. Bauman Stackexchange answer + Gitter discussion
-mutable struct VectorOfArray{T, N, A} <: AbstractVectorOfArray{T, N}
+mutable struct VectorOfArray{T, N, A} <: AbstractVectorOfArray{T, N, A}
   u::A # A <: AbstractVector{<: AbstractArray{T, N - 1}}
 end
 # VectorOfArray with an added series for time
-mutable struct DiffEqArray{T, N, A, B} <: AbstractDiffEqArray{T, N}
+mutable struct DiffEqArray{T, N, A, B} <: AbstractDiffEqArray{T, N, A}
   u::A # A <: AbstractVector{<: AbstractArray{T, N - 1}}
   t::B
 end
 
+Base.Array(VA::AbstractVectorOfArray{T,N,A}) where {T,N,A <: AbstractVector{<:AbstractVector}} = reduce(hcat,VA.u)
+function Base.Array(VA::AbstractVectorOfArray)
+  vecs = vec.(VA.u)
+  Array(reshape(reduce(hcat,vecs),size(VA.u[1])...,length(VA.u)))
+end
+
 VectorOfArray(vec::AbstractVector{T}, dims::NTuple{N}) where {T, N} = VectorOfArray{eltype(T), N, typeof(vec)}(vec)
 # Assume that the first element is representative all all other elements
 VectorOfArray(vec::AbstractVector) = VectorOfArray(vec, (size(vec[1])..., length(vec)))

test/basic_indexing.jl

@@ -6,10 +6,13 @@ testa = cat(recs..., dims=2)
 testva = VectorOfArray(recs)
 t = [1,2,3]
 diffeq = DiffEqArray(recs,t)
+@test Array(testva) == [1 4 7
+                        2 5 8
+                        3 6 9]
 
-testa[1:2, 1:2] == [1 4; 2 5]
-testva[1:2, 1:2] == [1 4; 2 5]
-testa[1:2, 1:2] == [1 4; 2 5]
+@test testa[1:2, 1:2] == [1 4; 2 5]
+@test testva[1:2, 1:2] == [1 4; 2 5]
+@test testa[1:2, 1:2] == [1 4; 2 5]
 
 # # ndims == 2
 recs = [rand(8) for i in 1:10]
@@ -73,3 +76,7 @@ testva[1:2, 1:2]
 # Test broadcast
 a = testva .+ rand(3,3)
 @test_broken a.= testva
+
+recs = [rand(2,2) for i in 1:5]
+testva = VectorOfArray(recs)
+@test Array(testva) isa Array{Float64,3}

test/interface_tests.jl

@@ -39,9 +39,9 @@ testva = VectorOfArray(recs)
 testa = cat(recs...,dims=3)
 @test convert(Array,testva) == testa
 
-recs = [[1, 2, 3], [3 5; 6 7], [8, 9, 10, 11]]
+recs = [[1 2; 3 4], [3 5; 6 7], [8 9; 10 11]]
 testva = VectorOfArray(recs)
-@test size(convert(Array,testva)) == (3,3)
+@test size(convert(Array,testva)) == (2,2,3)
 
 # create similar VectorOfArray
 recs = [rand(6) for i = 1:4]