fix: ArrayPartition mapreduce type inference, add tests

Author: AayushSabharwal

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))

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: