Add rrules for ArrayOfSimilarArrays ctors

Author: oschulz

src/array_of_similar_arrays.jl

@@ -70,29 +70,40 @@ struct ArrayOfSimilarArrays{
     data::P
 
     function ArrayOfSimilarArrays{T,M,N}(flat_data::AbstractArray{U,L}) where {T,M,N,L,U}
-        size_inner, size_outer = split_tuple(size(flat_data), Val{M}())
         require_ndims(flat_data, _add_vals(Val{M}(), Val{N}()))
         conv_parent = _convert_elype(T, flat_data)
         P = typeof(conv_parent)
         new{T,M,N,L,P}(conv_parent)
     end
+end
 
-    function ArrayOfSimilarArrays{T,M}(flat_data::AbstractArray{U,L}) where {T,M,L,U}
-        size_inner, size_outer = split_tuple(size(flat_data), Val{M}())
-        N = length(size_outer)
-        conv_parent = _convert_elype(T, flat_data)
-        P = typeof(conv_parent)
-        new{T,M,N,L,P}(conv_parent)
-    end
+function ArrayOfSimilarArrays{T,M}(flat_data::AbstractArray{U,L}) where {T,M,L,U}
+    _, size_outer = split_tuple(size(flat_data), Val{M}())
+    N = length(size_outer)
+    ArrayOfSimilarArrays{T,M,N}(flat_data)
 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)
 

test/array_of_similar_arrays.jl

@@ -6,6 +6,7 @@ using Test
 using ElasticArrays
 using UnsafeArrays
 using Adapt
+using ChainRulesTestUtils
 
 using Statistics
 using StatsBase: cov2cor
@@ -86,6 +87,8 @@ using StatsBase: cov2cor
         test_from_flat(VectorOfSimilarArrays{Float32}, VectorOfSimilarArrays{Float32,2,3,Array{Float32,3}}, Val(3))
         test_from_flat(VectorOfSimilarVectors{Float32}, VectorOfSimilarVectors{Float32,Array{Float32,2}}, Val(2))
         test_from_flat(VectorOfSimilarVectors{Float32}, VectorOfSimilarVectors{Float32,Array{Float32,2}}, Val(2))
+
+        test_rrule(ArrayOfSimilarArrays{Float64,2,2}, rand(2,3,4,5))
     end
 
 
@@ -115,6 +118,8 @@ using StatsBase: cov2cor
 
         r = @inferred(rand(5,5))
         @test @inferred(flatview(ArrayOfSimilarVectors(r))) == r
+
+        test_rrule(ArrayOfSimilarArrays{Float64,2,2}, [rand(2,3) for i in 1:5, j in 1:6])
     end
 
     @testset "add remove" begin