Use package extensions

Author: oschulz

Project.toml

@@ -8,6 +8,16 @@ ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
+[weakdeps]
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
+
+[extensions]
+ArraysOfArraysAdaptExt = "Adapt"
+ArraysOfArraysChainRulesCoreExt = "ChainRulesCore"
+ArraysOfArraysStaticArraysCoreExt = "StaticArraysCore"
+
 [compat]
 Adapt = "1, 2, 3, 4"
 ChainRulesCore = "1"
@@ -16,11 +26,12 @@ Statistics = "1"
 julia = "1.6"
 
 [extras]
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 ElasticArrays = "fdbdab4c-e67f-52f5-8c3f-e7b388dad3d4"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["ChainRulesTestUtils", "ElasticArrays", "StaticArrays", "StatsBase", "Test"]
+test = ["Adapt", "ChainRulesTestUtils", "ElasticArrays", "StaticArrays", "StatsBase", "Test"]

ext/ArraysOfArraysAdaptExt.jl

@@ -0,0 +1,28 @@
+# This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT).
+
+module ArraysOfArraysAdaptExt
+
+import Adapt
+using Adapt: adapt
+
+using ArraysOfArrays: ArrayOfSimilarArrays, VectorOfArrays
+using ArraysOfArrays: no_consistency_checks
+
+
+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 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
+
+
+end # module ArraysOfArraysAdaptExt

ext/ArraysOfArraysChainRulesCoreExt.jl

@@ -0,0 +1,37 @@
+# This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT).
+
+module ArraysOfArraysChainRulesCoreExt
+
+import ChainRulesCore
+using ChainRulesCore: NoTangent, unthunk
+
+using ArraysOfArrays: ArrayOfSimilarArrays
+using ArraysOfArrays: flatview
+
+
+function _aosa_ctor_fromflat_pullback(ΔΩ)
+    NoTangent(), flatview(convert(ArrayOfSimilarArrays, unthunk(ΔΩ)))
+end
+
+function ChainRulesCore.rrule(::Type{ArrayOfSimilarArrays{T,M,N}}, flat_data::AbstractArray{U,L}) where {T,M,N,L,U}
+    return ArrayOfSimilarArrays{T,M,N}(flat_data), _aosa_ctor_fromflat_pullback
+end
+
+_aosa_ctor_fromnested_pullback(ΔΩ) = NoTangent(), ΔΩ
+
+function ChainRulesCore.rrule(::Type{ArrayOfSimilarArrays{T,M,N}}, A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U}
+    return ArrayOfSimilarArrays{T,M,N}(A), _aosa_ctor_fromnested_pullback
+end
+
+
+function ChainRulesCore.rrule(::typeof(flatview), A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N}
+    function flatview_pullback(ΔΩ)
+        data = unthunk(ΔΩ)
+        NoTangent(), ArrayOfSimilarArrays{eltype(data),M,N}(data)
+    end
+    
+    return flatview(A), flatview_pullback
+end
+
+
+end # module ArraysOfArraysChainRulesCoreExt

ext/ArraysOfArraysStaticArraysCoreExt.jl

@@ -0,0 +1,26 @@
+# This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT).
+
+module ArraysOfArraysStaticArraysCoreExt
+
+import StaticArraysCore
+using StaticArraysCore: StaticArray, SVector
+
+import ArraysOfArrays
+using ArraysOfArrays: nestedview
+
+
+@inline ArraysOfArrays.flatview(A::AbstractArray{SA,N}) where {S,T,M,N,SA<:StaticArray{S,T,M}} =
+    reshape(reinterpret(T, A), size(SA)..., size(A)...)
+
+
+@inline function ArraysOfArrays.nestedview(A::AbstractArray{T}, SA::Type{SVector{S,T}}) where {T,S}
+    size_A = size(A)
+    size_A[1] == S || throw(DimensionMismatch("Length $S of static vector type does not match first dimension of array of size $size_A"))
+    reshape(reinterpret(SA, A), ArraysOfArrays._tail(size_A)...)
+end
+
+@inline ArraysOfArrays.nestedview(A::AbstractArray{T}, ::Type{SVector{S}}) where {T,S} =
+    nestedview(A, SVector{S,T})
+
+
+end # module ArraysOfArraysStaticArraysCoreExt

src/ArraysOfArrays.jl

@@ -13,16 +13,17 @@ ArraysOfArrays provides two different types of nested arrays:
 """
 module ArraysOfArrays
 
-using Adapt
 using Statistics
-using ChainRulesCore
-
-import StaticArraysCore
 
 include("util.jl")
 include("functions.jl")
 include("array_of_similar_arrays.jl")
 include("vector_of_arrays.jl")
-include("arrays_of_static_arrays.jl")
+
+@static if !isdefined(Base, :get_extension)
+    include("../ext/ArraysOfArraysAdaptExt.jl")
+    include("../ext/ArraysOfArraysChainRulesCoreExt.jl")
+    include("../ext/ArraysOfArraysStaticArraysCoreExt.jl")
+end
 
 end # module

src/array_of_similar_arrays.jl

@@ -85,25 +85,11 @@ end
 
 export ArrayOfSimilarArrays
 
-function _aosa_ctor_fromflat_pullback(ΔΩ)
-    NoTangent(), flatview(convert(ArrayOfSimilarArrays, unthunk(ΔΩ)))
-end
-
-function ChainRulesCore.rrule(::Type{ArrayOfSimilarArrays{T,M,N}}, flat_data::AbstractArray{U,L}) where {T,M,N,L,U}
-    return ArrayOfSimilarArrays{T,M,N}(flat_data), _aosa_ctor_fromflat_pullback
-end
-
 function ArrayOfSimilarArrays{T,M,N}(A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U}
     B = ArrayOfSimilarArrays{T,M,N}(Array{T}(undef, innersize(A)..., size(A)...))
     copyto!(B, A)
 end
 
-_aosa_ctor_fromnested_pullback(ΔΩ) = NoTangent(), ΔΩ
-
-function ChainRulesCore.rrule(::Type{ArrayOfSimilarArrays{T,M,N}}, A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U}
-    return ArrayOfSimilarArrays{T,M,N}(A), _aosa_ctor_fromnested_pullback
-end
-
 ArrayOfSimilarArrays{T}(A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} =
     ArrayOfSimilarArrays{T,M,N}(A)
 
@@ -143,15 +129,6 @@ the result may be freely changed without breaking the inner consistency of
 """
 flatview(A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} = A.data
 
-function ChainRulesCore.rrule(::typeof(flatview), A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N}
-    function flatview_pullback(ΔΩ)
-        data = unthunk(ΔΩ)
-        NoTangent(), ArrayOfSimilarArrays{eltype(data),M,N}(data)
-    end
-    
-    return flatview(A), flatview_pullback
-end
-
 
 Base.size(A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} = split_tuple(size(A.data), Val{M}())[2]
 
@@ -216,13 +193,6 @@ end
 Base.prepend!(dest::ArrayOfSimilarArrays{T,M,N}, src::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} =
     prepend!(dest, ArrayOfSimilarArrays(src))
 
-
-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/arrays_of_static_arrays.jl

@@ -398,16 +398,6 @@ function Base.empty!(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)