fix: ArrayPartition arithmetic type-stability

Author: AayushSabharwal

src/array_partition.jl

@@ -153,11 +153,11 @@ for op in (:*, :/)
 end
 
 function Base.:*(A::Number, B::ArrayPartition)
-    ArrayPartition(map(y -> Base.broadcast(*, A, y), B.x))
+    ArrayPartition(map(y -> A .* y, B.x))
 end
 
 function Base.:\(A::Number, B::ArrayPartition)
-    ArrayPartition(map(y -> Base.broadcast(/, y, A), B.x))
+    B / A
 end
 
 Base.:(==)(A::ArrayPartition, B::ArrayPartition) = A.x == B.x
@@ -284,7 +284,7 @@ recursive_eltype(A::ArrayPartition) = recursive_eltype(first(A.x))
 Base.iterate(A::ArrayPartition) = iterate(Chain(A.x))
 Base.iterate(A::ArrayPartition, state) = iterate(Chain(A.x), state)
 
-Base.length(A::ArrayPartition) = sum((length(x) for x in A.x))
+Base.length(A::ArrayPartition) = sum(broadcast(length, A.x))
 Base.size(A::ArrayPartition) = (length(A),)
 
 # redefine first and last to avoid slow and not type-stable indexing
@@ -323,21 +323,12 @@ function Broadcast.BroadcastStyle(::ArrayPartitionStyle,
     Broadcast.DefaultArrayStyle{N}()
 end
 
-combine_styles(args::Tuple{}) = Broadcast.DefaultArrayStyle{0}()
-@inline function combine_styles(args::Tuple{Any})
-    Broadcast.result_style(Broadcast.BroadcastStyle(args[1]))
-end
-@inline function combine_styles(args::Tuple{Any, Any})
-    Broadcast.result_style(Broadcast.BroadcastStyle(args[1]),
-        Broadcast.BroadcastStyle(args[2]))
-end
-@inline function combine_styles(args::Tuple)
-    Broadcast.result_style(Broadcast.BroadcastStyle(args[1]),
-        combine_styles(Base.tail(args)))
-end
+combine_styles(::Type{Tuple{}}) = Broadcast.DefaultArrayStyle{0}()
+combine_styles(::Type{T}) where {T} = Broadcast.result_style(Broadcast.BroadcastStyle(T.parameters[1]), combine_styles(Tuple{Base.tail((T.parameters...,))...}))
+
 
 function Broadcast.BroadcastStyle(::Type{ArrayPartition{T, S}}) where {T, S}
-    Style = combine_styles((S.parameters...,))
+    Style = combine_styles(S)
     ArrayPartitionStyle(Style)
 end
 

test/interface_tests.jl

@@ -110,7 +110,7 @@ end
 @test stack(testva; dims = 1) == [1 2 3; 4 5 6; 7 8 9]
 
 testva = VectorOfArray([VectorOfArray([ones(2,2), 2ones(2, 2)]), 3ones(2, 2, 2)])
-@test stack(testva) == [1.0 1.0; 1.0 1.0;;; 2.0 2.0; 2.0 2.0;;;; 3.0 3.0; 3.0;;; 3.0 3.0; 3.0 3.0]
+@test stack(testva) == [1.0 1.0; 1.0 1.0;;; 2.0 2.0; 2.0 2.0;;;; 3.0 3.0; 3.0 3.0;;; 3.0 3.0; 3.0 3.0]
 
 # convert array from VectorOfArray/DiffEqArray
 t = 1:8

test/partitions_test.jl

@@ -58,6 +58,7 @@ copyto!(p, c)
 ## inference tests
 
 x = ArrayPartition([1, 2], [3.0, 4.0])
+y = ArrayPartition(ArrayPartition([1], [2.0]), ArrayPartition([3], [4.0]))
 @test x[:, 1] == (1, 3.0)
 
 # similar partitions
@@ -69,6 +70,13 @@ x = ArrayPartition([1, 2], [3.0, 4.0])
 @test (@inferred similar(x, Int, (2, 2))) isa AbstractMatrix{Int}
 # @inferred similar(x, Int, Float64)
 
+@inferred similar(y)
+@test similar(y, (4,)) isa ArrayPartition{Float64}
+@test (@inferred similar(y, (2, 2))) isa AbstractMatrix{Float64}
+@inferred similar(y, Int)
+@test similar(y, Int, (4,)) isa ArrayPartition{Int}
+@test (@inferred similar(y, Int, (2, 2))) isa AbstractMatrix{Int}
+
 # Copy
 @inferred copy(x)
 @inferred copy(ArrayPartition(x, x))
@@ -95,6 +103,17 @@ x = ArrayPartition([1, 2], [3.0, 4.0])
 @inferred x + x
 @inferred x - x
 
+@inferred y + 5
+@inferred 5 + y
+@inferred y - 5
+@inferred 5 - y
+@inferred y * 5
+@inferred 5 * y
+@inferred y / 5
+@inferred 5 \ y
+@inferred y + y
+@inferred y - y
+
 # indexing
 @inferred first(x)
 @inferred last(x)
@@ -118,6 +137,7 @@ _scalar_op(y) = y + 1
 _broadcast_wrapper(y) = _scalar_op.(y)
 # Issue #8
 @inferred _broadcast_wrapper(x)
+@test_broken @inferred _broadcast_wrapper(y)
 
 # Testing map
 @test map(x -> x^2, x) == ArrayPartition(x.x[1] .^ 2, x.x[2] .^ 2)
@@ -151,11 +171,11 @@ function foo(y, x)
 end
 foo(xcde0, xce0)
 #@test 0 == @allocated foo(xcde0, xce0)
-function foo(y, x)
+function foo2(y, x)
     y .= y .+ 2 .* x
     nothing
 end
-foo(xcde0, xce0)
+foo2(xcde0, xce0)
 #@test 0 == @allocated foo(xcde0, xce0)
 
 # Custom AbstractArray types broadcasting