From 873a072c5b9cf7077280742a9939a9bd66061df7 Mon Sep 17 00:00:00 2001 From: Aayush Sabharwal Date: Fri, 5 Jan 2024 17:41:44 +0530 Subject: [PATCH 1/3] fix: ArrayPartition mapreduce type inference, add tests --- src/array_partition.jl | 4 ++-- test/partitions_test.jl | 8 ++++++++ 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/src/array_partition.jl b/src/array_partition.jl index 7506c92e..5be01442 100644 --- a/src/array_partition.jl +++ b/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)) diff --git a/test/partitions_test.jl b/test/partitions_test.jl index 87bd334a..9765a115 100644 --- a/test/partitions_test.jl +++ b/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: From 1ecc9661e3e34c7c7f864d967c9db876207047d3 Mon Sep 17 00:00:00 2001 From: Aayush Sabharwal Date: Fri, 5 Jan 2024 17:42:44 +0530 Subject: [PATCH 2/3] fix: VectorOfArray mapreduce, sum/prod performance, minor bugs --- src/vector_of_array.jl | 44 ++++++++++++++++++++++++++++++++--------- test/interface_tests.jl | 38 +++++++++++++++++++++++++++++++++++ 2 files changed, 73 insertions(+), 9 deletions(-) diff --git a/src/vector_of_array.jl b/src/vector_of_array.jl index 9dcfc409..14476500 100644 --- a/src/vector_of_array.jl +++ b/src/vector_of_array.jl @@ -133,7 +133,16 @@ function VectorOfArray(vec::AbstractVector{T}, ::NTuple{N}) where {T, N} 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,6 +491,10 @@ function Base.append!(VA::AbstractVectorOfArray{T, N}, 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 @@ -489,14 +502,19 @@ function Base.view(A::AbstractVectorOfArray, I::Vararg{Any,M}) where {M} @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 @@ function Base.convert(::Type{Array}, VA::AbstractVectorOfArray) 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 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 diff --git a/test/interface_tests.jl b/test/interface_tests.jl index d2f4030b..3aba249a 100644 --- a/test/interface_tests.jl +++ b/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] From 1f9b57774bda4f96121f5a226e40941e3da1b73a Mon Sep 17 00:00:00 2001 From: Aayush Sabharwal Date: Fri, 5 Jan 2024 17:43:07 +0530 Subject: [PATCH 3/3] feat: add and test adjoints for broadcast arithmetic on VoA --- ext/RecursiveArrayToolsZygoteExt.jl | 112 +++++++++++++++++++++++++++- test/adjoints.jl | 23 +++++- 2 files changed, 132 insertions(+), 3 deletions(-) diff --git a/ext/RecursiveArrayToolsZygoteExt.jl b/ext/RecursiveArrayToolsZygoteExt.jl index ead05de1..ac2f1622 100644 --- a/ext/RecursiveArrayToolsZygoteExt.jl +++ b/ext/RecursiveArrayToolsZygoteExt.jl @@ -100,6 +100,11 @@ end 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 @@ end 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̄) diff --git a/test/adjoints.jl b/test/adjoints.jl index e06035af..c657dcf9 100644 --- a/test/adjoints.jl +++ b/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)