Merge pull request #202 from SciML/gpu

add GPUArrays and GPU differentiation test

Author: ChrisRackauckas

Project.toml

@@ -9,6 +9,7 @@ ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
+GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
@@ -21,6 +22,7 @@ ArrayInterface = "2.7, 3.0, 4, 5"
 ChainRulesCore = "0.10.7, 1"
 DocStringExtensions = "0.8"
 FillArrays = "0.11, 0.12, 0.13"
+GPUArrays = "8"
 RecipesBase = "0.7, 0.8, 1.0"
 StaticArrays = "0.12, 1.0"
 ZygoteRules = "0.2"

src/RecursiveArrayTools.jl

@@ -24,6 +24,10 @@ include("zygote.jl")
 
 Base.show(io::IO, x::Union{ArrayPartition,AbstractVectorOfArray}) = invoke(show, Tuple{typeof(io), Any}, io, x)
 
+import GPUArrays
+Base.convert(T::Type{<:GPUArrays.AbstractGPUArray}, VA::AbstractVectorOfArray) = T(VA)
+ChainRulesCore.rrule(T::Type{<:GPUArrays.AbstractGPUArray}, xs::AbstractVectorOfArray) = T(xs), ȳ -> (NoTangent(),ȳ)
+
 export VectorOfArray, DiffEqArray, AbstractVectorOfArray, AbstractDiffEqArray,
        AllObserved, vecarr_to_arr, vecarr_to_vectors, tuples
 

test/gpu/vectorofarray_gpu.jl

@@ -1,14 +1,14 @@
-using RecursiveArrayTools, CUDA, Test
+using RecursiveArrayTools, CUDA, Test, Zygote
 CUDA.allowscalar(false)
 
 # Test indexing with colon
 x = zeros(5)
-y = VectorOfArray([x,x,x])
-y[:,:]
+y = VectorOfArray([x, x, x])
+y[:, :]
 
 x = CUDA.zeros(5)
-y = VectorOfArray([x,x,x])
-y[:,:]
+y = VectorOfArray([x, x, x])
+y[:, :]
 
 # Test indexing with boolean masks and colon
 nx, ny, nt = 3, 4, 5
@@ -22,3 +22,12 @@ va = VectorOfArray([slice for slice in eachslice(x, dims=3)])
 xc = Array(x)
 mc = Array(m)
 @test xc[mc, :] ≈ Array(va[m, :])
+
+# Check differentiation with GPUs
+
+p = cu([1.0, 2.0])
+function f(p)
+    x = VectorOfArray([p, p])
+    sum(CuArray(x))
+end
+Zygote.gradient(f, p)