diff --git a/Project.toml b/Project.toml index 7adf8711..22d07c4f 100644 --- a/Project.toml +++ b/Project.toml @@ -6,8 +6,6 @@ version = "2.3.0" [deps] ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" -NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56" -OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01" Requires = "ae029012-a4dd-5104-9daa-d747884805df" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" @@ -23,8 +21,10 @@ ZygoteRules = "0.2" julia = "1.3" [extras] +NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56" +OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d" [targets] -test = ["Test", "Unitful"] +test = ["NLsolve", "OrdinaryDiffEq", "Test", "Unitful"] diff --git a/src/array_partition.jl b/src/array_partition.jl index e4ec24d3..65743356 100644 --- a/src/array_partition.jl +++ b/src/array_partition.jl @@ -230,15 +230,15 @@ ArrayPartitionStyle(::S, ::Val{N}) where {S,N} = ArrayPartitionStyle(S(Val(N))) ArrayPartitionStyle(::Val{N}) where N = ArrayPartitionStyle{Broadcast.DefaultArrayStyle{N}}() # promotion rules -function Broadcast.BroadcastStyle(::ArrayPartitionStyle{AStyle}, ::ArrayPartitionStyle{BStyle}) where {AStyle, BStyle} +@inline function Broadcast.BroadcastStyle(::ArrayPartitionStyle{AStyle}, ::ArrayPartitionStyle{BStyle}) where {AStyle, BStyle} ArrayPartitionStyle(Broadcast.BroadcastStyle(AStyle(), BStyle())) end -Broadcast.BroadcastStyle(::ArrayPartitionStyle, ::Broadcast.DefaultArrayStyle{0}) = Broadcast.DefaultArrayStyle{1}() +Broadcast.BroadcastStyle(::ArrayPartitionStyle{Style}, ::Broadcast.DefaultArrayStyle{0}) where Style<:Broadcast.BroadcastStyle = ArrayPartitionStyle{Style}() Broadcast.BroadcastStyle(::ArrayPartitionStyle, ::Broadcast.DefaultArrayStyle{N}) where N = Broadcast.DefaultArrayStyle{N}() combine_styles(args::Tuple{}) = Broadcast.DefaultArrayStyle{0}() -combine_styles(args::Tuple{Any}) = Broadcast.result_style(Broadcast.BroadcastStyle(args[1])) -combine_styles(args::Tuple{Any, Any}) = Broadcast.result_style(Broadcast.BroadcastStyle(args[1]), Broadcast.BroadcastStyle(args[2])) +@inline combine_styles(args::Tuple{Any}) = Broadcast.result_style(Broadcast.BroadcastStyle(args[1])) +@inline combine_styles(args::Tuple{Any, Any}) = Broadcast.result_style(Broadcast.BroadcastStyle(args[1]), Broadcast.BroadcastStyle(args[2])) @inline combine_styles(args::Tuple) = Broadcast.result_style(Broadcast.BroadcastStyle(args[1]), combine_styles(Base.tail(args))) function Broadcast.BroadcastStyle(::Type{ArrayPartition{T,S}}) where {T, S} @@ -256,7 +256,7 @@ end @inline function Base.copyto!(dest::ArrayPartition, bc::Broadcast.Broadcasted{ArrayPartitionStyle{Style}}) where Style N = npartitions(dest, bc) - for i in 1:N + @inbounds for i in 1:N copyto!(dest.x[i], unpack(bc, i)) end dest diff --git a/test/partitions_test.jl b/test/partitions_test.jl index b032d05a..8caad927 100644 --- a/test/partitions_test.jl +++ b/test/partitions_test.jl @@ -119,3 +119,19 @@ for sizes in S @test all([x[i] == y[i] for i in eachindex(x)]) @test all([x[i] == y_array[i] for i in eachindex(x)]) end + +# Non-allocating broadcast +xce0 = ArrayPartition(zeros(2),[0.]) +xcde0 = copy(xce0) +function foo(y, x) + y .= y .+ x + nothing +end +foo(xcde0, xce0) +#@test 0 == @allocated foo(xcde0, xce0) +function foo(y, x) + y .= y .+ 2 .* x + nothing +end +foo(xcde0, xce0) +#@test 0 == @allocated foo(xcde0, xce0)