fix broadcast against non APs and test upstream

Author: ChrisRackauckas

src/array_partition.jl

@@ -233,6 +233,8 @@ ArrayPartitionStyle(::Val{N}) where N = ArrayPartitionStyle{Broadcast.DefaultArr
 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, ::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]))
@@ -252,7 +254,7 @@ end
     ArrayPartition(f, N)
 end
 
-@inline function Base.copyto!(dest::ArrayPartition, bc::Broadcast.Broadcasted)
+@inline function Base.copyto!(dest::ArrayPartition, bc::Broadcast.Broadcasted{ArrayPartitionStyle{Style}}) where Style
     N = npartitions(dest, bc)
     for i in 1:N
         copyto!(dest.x[i], unpack(bc, i))

test/partitions_test.jl

@@ -26,7 +26,21 @@ a = 5
 @. p = p*p2
 K = p.*p2
 
-@test_broken p.*rand(10)
+x = rand(10)
+y = p.*x
+@test y[1:5] == p.x[1] .* x[1:5]
+@test y[6:10] == p.x[2] .* x[6:10]
+y = p.*x'
+for i in 1:10
+  @test y[1:5,i] == p.x[1] .* x[i]
+  @test y[6:10,i] == p.x[2] .* x[i]
+end
+y = p .* p'
+@test y[1:5,1:5] == p.x[1] .* p.x[1]'
+@test y[6:10,6:10] == p.x[2] .* p.x[2]'
+@test y[1:5,6:10] == p.x[1] .* p.x[2]'
+@test y[6:10,1:5] == p.x[2] .* p.x[1]'
+
 b = rand(10)
 c = rand(10)
 copyto!(b,p)
@@ -94,13 +108,11 @@ S = [
     ]
 
 for sizes in S
-  x = ArrayPartition( randn.(sizes[1]) ) 
+  x = ArrayPartition( randn.(sizes[1]) )
   y = ArrayPartition( zeros.(sizes[2]) )
   y_array = zeros(length(x))
   copyto!(y,x)           #testing Base.copyto!(dest::ArrayPartition,A::ArrayPartition)
   copyto!(y_array,x)     #testing Base.copyto!(dest::Array,A::ArrayPartition)
   @test all([x[i] == y[i] for i in eachindex(x)])
   @test all([x[i] == y_array[i] for i in eachindex(x)])
 end
-  
-

test/upstream.jl

@@ -8,9 +8,9 @@ u0 = ArrayPartition([1.0,0.0],[0.0])
 @test ArrayInterface.zeromatrix(u0) isa Matrix
 tspan = (0.0,100.0)
 prob = ODEProblem(lorenz,u0,tspan)
-sol = solve(prob,Tsit5());
-sol = solve(prob,AutoTsit5(Rosenbrock23(autodiff=false)));
-@test_broken sol = solve(prob,AutoTsit5(Rosenbrock23()));
+sol = solve(prob,Tsit5())
+sol = solve(prob,AutoTsit5(Rosenbrock23(autodiff=false)))
+sol = solve(prob,AutoTsit5(Rosenbrock23()))
 
 function f!(F, vars)
     x = vars.x[1]
@@ -24,7 +24,7 @@ end
 # To show that the function works
 F = ArrayPartition([0.0 0.0],[0.0, 0.0])
 u0= ArrayPartition([0.1; 1.2], [0.1; 1.2])
-result = mymodel(F, u0)
+result = f!(F, u0)
 
 # To show the NLsolve error that results with ArrayPartitions:
 nlsolve(f!, ArrayPartition([0.1; 1.2], [0.1; 1.2]))