Add support for Adapt.jl

Author: oschulz

Project.toml

@@ -3,11 +3,13 @@ uuid = "65a8f2f4-9b39-5baf-92e2-a9cc46fdf018"
 version = "0.5.0"
 
 [deps]
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 UnsafeArrays = "c4a57d5a-5b31-53a6-b365-19f8c011fbd6"
 
 [compat]
+Adapt = "1"
 Requires = "0.5,1"
 UnsafeArrays = "1"
 julia = "1"

src/ArraysOfArrays.jl

@@ -4,6 +4,7 @@ __precompile__(true)
 
 module ArraysOfArrays
 
+using Adapt
 using Requires
 using Statistics
 using UnsafeArrays

src/array_of_similar_arrays.jl

@@ -201,6 +201,12 @@ UnsafeArrays.unsafe_uview(A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} =
     ArrayOfSimilarArrays{T,M,N}(uview(A.data))
 
 
+function Adapt.adapt_structure(to, A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N}
+    adapted_data = adapt(to, A.data)
+    ArrayOfSimilarArrays{eltype(adapted_data),M,N}(adapted_data)
+end
+
+
 function innermap(f::Base.Callable, A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N}
     new_data = map(f, A.data)
     U = eltype(new_data)

src/vector_of_arrays.jl

@@ -397,6 +397,16 @@ function UnsafeArrays.uview(A::VectorOfArrays)
 end
 
 
+function Adapt.adapt_structure(to, A::VectorOfArrays)
+    VectorOfArrays(
+        adapt(to, A.data),
+        adapt(to, A.elem_ptr),
+        adapt(to, A.kernel_size),
+        no_consistency_checks
+    )
+end
+
+
 function innermap(f::Base.Callable, A::VectorOfArrays)
     new_data = map(f, A.data)
     VectorOfArrays(new_data, A.elem_ptr, A.kernel_size, simple_consistency_checks)

test/array_of_similar_arrays.jl

@@ -5,6 +5,7 @@ using Test
 
 using ElasticArrays
 using UnsafeArrays
+using Adapt
 
 using Statistics
 using StatsBase: cov2cor
@@ -160,6 +161,14 @@ using StatsBase: cov2cor
     end
 
 
+    @testset "adapt" begin
+        A_flat = rand(2,3,4,5,6)
+        A_nested = nestedview(A_flat, 2)
+        @test @inferred(adapt(identity, A_nested)) == A_nested
+        @test typeof(adapt(identity, A_nested)) == typeof(A_nested)
+    end
+
+
     @testset "deepcopy" begin
         A = ArrayOfSimilarArrays{Float64,1}(rand_flat_array(Val(1)))
         @test (@inferred deepcopy(A)) == A
@@ -274,7 +283,6 @@ using StatsBase: cov2cor
         end
     end
 
-
     @testset "examples" begin
         A_flat = rand(2,3,4,5,6)
         A_nested = nestedview(A_flat, 2)
@@ -323,6 +331,7 @@ using StatsBase: cov2cor
         @test_throws ArgumentError pop!(A_nested)
 
     end
+
     @testset "misc" begin
         N = 4
         r1 = rand(1,4); r2 = rand(1,4); r3 = rand(1,4); r4 = rand(1,4)

test/vector_of_arrays.jl

@@ -6,6 +6,7 @@ using Statistics
 using Test
 
 using UnsafeArrays
+using Adapt
 
 using ArraysOfArrays: full_consistency_checks, append_elemptr!, element_ptr
 
@@ -222,6 +223,17 @@ using ArraysOfArrays: full_consistency_checks, append_elemptr!, element_ptr
     end
 
 
+    @testset "adapt" begin
+        A1 = VectorOfArrays(ref_AoA1(Float32, 3))
+        @test @inferred(adapt(identity, A1)) == A1
+        @test typeof(adapt(identity, A1)) == typeof(A1)
+
+        A3 = VectorOfArrays(ref_AoA3(Float32, 3))
+        @test @inferred(adapt(identity, A3)) == A3
+        @test typeof(adapt(identity, A3)) == typeof(A3)
+    end
+
+
     @testset "examples" begin
         VA = @inferred(VectorOfArrays{Float64, 2}())