fix: mapreduce type stability, VoA broadcast adjoints

ext/RecursiveArrayToolsZygoteExt.jl

@@ -100,6 +100,11 @@
     end
 end
 
+@adjoint function Base.copy(u::VectorOfArray)
+    copy(u),
+    y ->  (copy(y),)
+end
+
 @adjoint function DiffEqArray(u, t)
     DiffEqArray(u, t),
     y -> begin
@@ -117,19 +122,122 @@
     A.x, literal_ArrayPartition_x_adjoint
 end
 
-@adjoint function Array(VA::AbstractVectorOfArray)
+@adjoint function Base.Array(VA::AbstractVectorOfArray)
     Array(VA),
     y -> (Array(y),)
 end
 
+@adjoint function Base.view(A::AbstractVectorOfArray, I...)
+    view(A, I...),
+    y -> (view(y, I...), ntuple(_ -> nothing, length(I))...)
+end
 
 ChainRulesCore.ProjectTo(a::AbstractVectorOfArray) = ChainRulesCore.ProjectTo{VectorOfArray}((sz = size(a)))
 
-function (p::ChainRulesCore.ProjectTo{VectorOfArray})(x)
+function (p::ChainRulesCore.ProjectTo{VectorOfArray})(x::Union{AbstractArray,AbstractVectorOfArray})
     arr = reshape(x, p.sz)
     return VectorOfArray([arr[:, i] for i in 1:p.sz[end]])
 end
 
+@adjoint function Broadcast.broadcasted(::typeof(+), x::AbstractVectorOfArray, y::Union{Zygote.Numeric, AbstractVectorOfArray})
+    broadcast(+, x, y), ȳ -> (nothing, map(x -> Zygote.unbroadcast(x, ȳ), (x, y))...)
+end
+@adjoint function Broadcast.broadcasted(::typeof(+), x::Zygote.Numeric, y::AbstractVectorOfArray)
+    broadcast(+, x, y), ȳ -> (nothing, map(x -> Zygote.unbroadcast(x, ȳ), (x, y))...)
+end
+
+_minus(Δ) = .-Δ
+_minus(::Nothing) = nothing
+
+@adjoint function Broadcast.broadcasted(::typeof(-), x::AbstractVectorOfArray, y::Union{AbstractVectorOfArray, Zygote.Numeric})
+    x .- y, Δ -> (nothing, Zygote.unbroadcast(x, Δ), _minus(Zygote.unbroadcast(y, Δ)))
+end
+@adjoint function Broadcast.broadcasted(::typeof(*), x::AbstractVectorOfArray, y::Union{AbstractVectorOfArray, Zygote.Numeric})
+    (
+        x.*y,
+        Δ -> (nothing, Zygote.unbroadcast(x, Δ .* conj.(y)), Zygote.unbroadcast(y, Δ .* conj.(x)))
+    )
+end
+@adjoint function Broadcast.broadcasted(::typeof(/), x::AbstractVectorOfArray, y::Union{AbstractVectorOfArray, Zygote.Numeric})
+  res = x ./ y
+  res, Δ -> (nothing, Zygote.unbroadcast(x, Δ ./ conj.(y)), Zygote.unbroadcast(y, .-Δ .* conj.(res ./ y)))
+end
+@adjoint function Broadcast.broadcasted(::typeof(-), x::Zygote.Numeric, y::AbstractVectorOfArray)
+    x .- y, Δ -> (nothing, Zygote.unbroadcast(x, Δ), _minus(Zygote.unbroadcast(y, Δ)))
+end
+@adjoint function Broadcast.broadcasted(::typeof(*), x::Zygote.Numeric, y::AbstractVectorOfArray)
+    (
+        x.*y,
+        Δ -> (nothing, Zygote.unbroadcast(x, Δ .* conj.(y)), Zygote.unbroadcast(y, Δ .* conj.(x)))
+    )
+end
+@adjoint function Broadcast.broadcasted(::typeof(/), x::Zygote.Numeric, y::AbstractVectorOfArray)
+  res = x ./ y
+  res, Δ -> (nothing, Zygote.unbroadcast(x, Δ ./ conj.(y)), Zygote.unbroadcast(y, .-Δ .* conj.(res ./ y)))
+end
+@adjoint function Broadcast.broadcasted(::typeof(-), x::AbstractVectorOfArray)
+    .-x, Δ -> (nothing, _minus(Δ))
+end
+
+@adjoint function Broadcast.broadcasted(::typeof(Base.literal_pow), ::typeof(^), x::AbstractVectorOfArray, exp::Val{p}) where p
+  y = Base.literal_pow.(^, x, exp)
+  y, ȳ -> (nothing, nothing, ȳ .* p .* conj.(x .^ (p - 1)), nothing)
+end
+
+@adjoint Broadcast.broadcasted(::typeof(identity), x::AbstractVectorOfArray) = x, Δ -> (nothing, Δ)
+
+@adjoint function Broadcast.broadcasted(::typeof(tanh), x::AbstractVectorOfArray)
+  y = tanh.(x)
+  y, ȳ -> (nothing, ȳ .* conj.(1 .- y.^2))
+end
+
+@adjoint Broadcast.broadcasted(::typeof(conj), x::AbstractVectorOfArray) =
+  conj.(x), z̄ -> (nothing, conj.(z̄))
+
+@adjoint Broadcast.broadcasted(::typeof(real), x::AbstractVectorOfArray) =
+  real.(x), z̄ -> (nothing, real.(z̄))
+
+@adjoint Broadcast.broadcasted(::typeof(imag), x::AbstractVectorOfArray) =
+  imag.(x), z̄ -> (nothing, im .* real.(z̄))
+
+@adjoint Broadcast.broadcasted(::typeof(abs2), x::AbstractVectorOfArray) =
+  abs2.(x), z̄ -> (nothing, 2 .* real.(z̄) .* x)
+
+@adjoint function Broadcast.broadcasted(::typeof(+), a::AbstractVectorOfArray{<:Number}, b::Bool)
+  y = b === false ? a : a .+ b
+  y, Δ -> (nothing, Δ, nothing)
+end
+@adjoint function Broadcast.broadcasted(::typeof(+), b::Bool, a::AbstractVectorOfArray{<:Number})
+  y = b === false ? a : b .+ a
+  y, Δ -> (nothing, nothing, Δ)
+end
+
+@adjoint function Broadcast.broadcasted(::typeof(-), a::AbstractVectorOfArray{<:Number}, b::Bool)
+  y = b === false ? a : a .- b
+  y, Δ -> (nothing, Δ, nothing)
+end
+@adjoint function Broadcast.broadcasted(::typeof(-), b::Bool, a::AbstractVectorOfArray{<:Number})
+  b .- a, Δ -> (nothing, nothing, .-Δ)
+end
+
+@adjoint function Broadcast.broadcasted(::typeof(*), a::AbstractVectorOfArray{<:Number}, b::Bool)
+  if b === false
+    zero(a), Δ -> (nothing, zero(Δ), nothing)
+  else
+    a, Δ -> (nothing, Δ, nothing)
+  end
+end
+@adjoint function Broadcast.broadcasted(::typeof(*), b::Bool, a::AbstractVectorOfArray{<:Number})
+  if b === false
+    zero(a), Δ -> (nothing, nothing, zero(Δ))
+  else
+    a, Δ -> (nothing, nothing, Δ)
+  end
+end
+
+@adjoint Broadcast.broadcasted(::Type{T}, x::AbstractVectorOfArray) where {T<:Number} =
+  T.(x), ȳ -> (nothing, Zygote._project(x, ȳ),)
+
 function Zygote.unbroadcast(x::AbstractVectorOfArray, x̄)
     N = ndims(x̄)
     if length(x) == length(x̄)

src/array_partition.jl

@@ -165,8 +165,8 @@ Base.:(==)(A::ArrayPartition, B::ArrayPartition) = A.x == B.x
 ## Iterable Collection Constructs
 
 Base.map(f, A::ArrayPartition) = ArrayPartition(map(x -> map(f, x), A.x))
-function Base.mapreduce(f, op, A::ArrayPartition)
-    mapreduce(f, op, (mapreduce(f, op, x) for x in A.x))
+function Base.mapreduce(f, op, A::ArrayPartition{T}; kwargs...) where {T}
+    mapreduce(f, op, (i for i in A); kwargs...)
 end
 Base.filter(f, A::ArrayPartition) = ArrayPartition(map(x -> filter(f, x), A.x))
 Base.any(f, A::ArrayPartition) = any(f, (any(f, x) for x in A.x))

src/vector_of_array.jl

@@ -133,7 +133,16 @@
     VectorOfArray{eltype(T), N, typeof(vec)}(vec)
 end
 # Assume that the first element is representative of all other elements
-VectorOfArray(vec::AbstractVector) = VectorOfArray(vec, (size(vec[1])..., length(vec)))
+function VectorOfArray(vec::AbstractVector)
+    T = eltype(vec[1])
+    N = ndims(vec[1])
+    if all(x isa Union{<:AbstractArray, <:AbstractVectorOfArray} for x in vec)
+        A = Vector{Union{typeof.(vec)...}}
+    else
+        A = typeof(vec)
+    end
+    VectorOfArray{T, N + 1, A}(vec)
+end
 function VectorOfArray(vec::AbstractVector{VT}) where {T, N, VT <: AbstractArray{T, N}}
     VectorOfArray{T, N + 1, typeof(vec)}(vec)
 end
@@ -482,21 +491,30 @@
     return VA
 end
 
+function Base.stack(VA::AbstractVectorOfArray; dims = :)
+    stack(VA.u; dims)
+end
+
 # AbstractArray methods
 function Base.view(A::AbstractVectorOfArray, I::Vararg{Any,M}) where {M}
     @inline
     J = map(i->Base.unalias(A,i), to_indices(A, I))
     @boundscheck checkbounds(A, J...)
     SubArray(IndexStyle(A), A, J, Base.index_dimsum(J...))
 end
+function Base.SubArray(parent::AbstractVectorOfArray, indices::Tuple)
+    @inline
+    SubArray(IndexStyle(Base.viewindexing(indices), IndexStyle(parent)), parent, Base.ensure_indexable(indices), Base.index_dimsum(indices...))
+end
+Base.isassigned(VA::AbstractVectorOfArray, idxs...) = checkbounds(Bool, VA, idxs...)
 Base.check_parent_index_match(::RecursiveArrayTools.AbstractVectorOfArray{T,N}, ::NTuple{N,Bool}) where {T,N} = nothing
 Base.ndims(::AbstractVectorOfArray{T, N}) where {T, N} = N
 function Base.checkbounds(::Type{Bool}, VA::AbstractVectorOfArray, idx...)
     if checkbounds(Bool, VA.u, last(idx))
         if last(idx) isa Integer
-            return all(checkbounds.(Bool, (VA.u[last(idx)],), Base.front(idx)))
+            return all(checkbounds.(Bool, (VA.u[last(idx)],), Base.front(idx)...))
         else
-            return all(checkbounds.(Bool, VA.u[last(idx)], Base.front(idx)))
+            return all(checkbounds.(Bool, VA.u[last(idx)], tuple.(Base.front(idx))...))
         end
     end
     return false
@@ -595,10 +613,14 @@
 end
 
 # statistics
-@inline Base.sum(f, VA::AbstractVectorOfArray) = sum(f, Array(VA))
-@inline Base.sum(VA::AbstractVectorOfArray; kwargs...) = sum(Array(VA); kwargs...)
-@inline Base.prod(f, VA::AbstractVectorOfArray) = prod(f, Array(VA))
-@inline Base.prod(VA::AbstractVectorOfArray; kwargs...) = prod(Array(VA); kwargs...)
+@inline Base.sum(VA::AbstractVectorOfArray; kwargs...) = sum(identity, VA; kwargs...)
+@inline function Base.sum(f, VA::AbstractVectorOfArray; kwargs...)
+    mapreduce(f, Base.add_sum, VA; kwargs...)
+end
+@inline Base.prod(VA::AbstractVectorOfArray; kwargs...) = prod(identity, VA; kwargs...)
+@inline function Base.prod(f, VA::AbstractVectorOfArray; kwargs...)
+    mapreduce(f, Base.mul_prod, VA; kwargs...)
+end
 
 @inline Statistics.mean(VA::AbstractVectorOfArray; kwargs...) = mean(Array(VA); kwargs...)
 @inline function Statistics.median(VA::AbstractVectorOfArray; kwargs...)
@@ -638,8 +660,12 @@
 end
 
 Base.map(f, A::RecursiveArrayTools.AbstractVectorOfArray) = map(f, A.u)
-function Base.mapreduce(f, op, A::AbstractVectorOfArray)
-    mapreduce(f, op, (mapreduce(f, op, x) for x in A.u))
+
+function Base.mapreduce(f, op, A::AbstractVectorOfArray; kwargs...)
+    mapreduce(f, op, view(A, ntuple(_ -> :, ndims(A))...); kwargs...)
+end
+function Base.mapreduce(f, op, A::AbstractVectorOfArray{T,1,<:AbstractVector{T}}; kwargs...) where {T}
+    mapreduce(f, op, A.u; kwargs...)
 end
 
 ## broadcasting

test/adjoints.jl

@@ -39,7 +39,27 @@ end
 
 function loss7(x)
     _x = VectorOfArray([x .* i for i in 1:5])
-    return sum(abs2, x .- 1)
+    return sum(abs2, _x .- 1)
+end
+
+# use a bunch of broadcasts to test all the adjoints
+function loss8(x)
+    _x = VectorOfArray([x .* i for i in 1:5])
+    res = copy(_x)
+    res = res .+ _x
+    res = res .+ 1
+    res = res .* _x
+    res = res .* 2.0
+    res = res .* res
+    res = res ./ 2.0
+    res = res ./ _x
+    res = 3.0 .- res
+    res = .-res
+    res = identity.(Base.literal_pow.(^, res, Val(2)))
+    res = tanh.(res)
+    res = res .+ im .* res
+    res = conj.(res) .+ real.(res) .+ imag.(res) .+ abs2.(res)
+    return sum(abs2, res)
 end
 
 x = float.(6:10)
@@ -51,3 +71,4 @@ loss(x)
 @test Zygote.gradient(loss5, x)[1] == ForwardDiff.gradient(loss5, x)
 @test Zygote.gradient(loss6, x)[1] == ForwardDiff.gradient(loss6, x)
 @test Zygote.gradient(loss7, x)[1] == ForwardDiff.gradient(loss7, x)
+@test Zygote.gradient(loss8, x)[1] == ForwardDiff.gradient(loss8, x)

test/interface_tests.jl

@@ -57,6 +57,44 @@ push!(testda, [-1, -2, -3, -4])
 @test_throws MethodError push!(testda, [-1 -2 -3 -4])
 @test_throws MethodError push!(testda, [-1 -2; -3 -4])
 
+# Type inference
+@inferred sum(testva)
+@inferred sum(VectorOfArray([VectorOfArray([zeros(4,4)])]))
+@inferred mapreduce(string, *, testva)
+
+# mapreduce
+testva = VectorOfArray([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
+@test mapreduce(x -> string(x) * "q", *, testva) == "1q2q3q4q5q6q7q8q9q"
+
+testvb = VectorOfArray([rand(1:10, 3, 3, 3) for _ in 1:4])
+arrvb = Array(testvb)
+for i in 1:ndims(arrvb)
+    @test sum(arrvb; dims=i) == sum(testvb; dims=i)
+    @test prod(arrvb; dims=i) == prod(testvb; dims=i)
+    @test mapreduce(string, *, arrvb; dims=i) == mapreduce(string, *, testvb; dims=i)
+end
+
+# Test when ndims == 1
+testvb = VectorOfArray(collect(1.0:0.1:2.0))
+arrvb = Array(testvb)
+@test sum(arrvb) == sum(testvb)
+@test prod(arrvb) == prod(testvb)
+@test mapreduce(string, *, arrvb) == mapreduce(string, *, testvb)
+
+# view
+testvc = VectorOfArray([rand(1:10, 3, 3) for _ in 1:3])
+arrvc = Array(testvc)
+for idxs in [(2, 2, :), (2, :, 2), (:, 2, 2), (:, :, 2), (:, 2, :), (2, : ,:), (:, :, :)]
+    arr_view = view(arrvc, idxs...)
+    voa_view = view(testvc, idxs...)
+    @test size(arr_view) == size(voa_view)
+    @test all(arr_view .== voa_view)
+end
+
+# test stack
+@test stack(testva) == [1 4 7; 2 5 8; 3 6 9]
+@test stack(testva; dims = 1) == [1 2 3; 4 5 6; 7 8 9]
+
 # convert array from VectorOfArray/DiffEqArray
 t = 1:8
 recs = [rand(10, 7) for i in 1:8]

test/partitions_test.jl

@@ -104,6 +104,14 @@ x = ArrayPartition([1, 2], [3.0, 4.0])
 @inferred recursive_one(x)
 @inferred recursive_bottom_eltype(x)
 
+# mapreduce
+@inferred Union{Int, Float64} sum(x)
+@inferred sum(ArrayPartition(ArrayPartition(zeros(4,4))))
+@inferred sum(ArrayPartition(ArrayPartition(zeros(4))))
+@inferred sum(ArrayPartition(zeros(4,4)))
+@inferred mapreduce(string, *, x)
+@test mapreduce(i -> string(i) * "q", *, x) == "1q2q3.0q4.0q"
+
 # broadcasting
 _scalar_op(y) = y + 1
 # Can't do `@inferred(_scalar_op.(x))` so we wrap that in a function: