Fix ArrayPartition A[:] DimensionMismatch (#661)

src/array_partition.jl

@@ -317,12 +317,14 @@ Base.@propagate_inbounds function Base.getindex(A::ArrayPartition, i::Int)
 end
 
 """
-    getindex(A::ArrayPartition, i::Colon, j...)
+    getindex(A::ArrayPartition, i::Colon, j, js...)
 
-Returns the entry at index `j...` of  every partition of `A`.
+Returns the entry at index `j, js...` of every partition of `A`.
+Requires at least one index after the colon so that `A[:]` keeps the
+flattening `getindex(::ArrayPartition, ::Colon)` method.
 """
-Base.@propagate_inbounds function Base.getindex(A::ArrayPartition, i::Colon, j...)
-    return getindex.(A.x, (j...,))
+Base.@propagate_inbounds function Base.getindex(A::ArrayPartition, i::Colon, j, js...)
+    return getindex.(A.x, (j, js...))
 end
 
 """

test/Core/partitions_test.jl

@@ -90,6 +90,30 @@ x = AP[[1, 2], [3.0, 4.0]]
 y = ArrayPartition(AP[[1], [2.0]], AP[[3], [4.0]])
 @test x[:, 1] == (1, 3.0)
 
+# Issue #661: A[:] must flatten, not hit the (colon, j...) broadcast path
+@testset "ArrayPartition A[:] flattening (issue #661)" begin
+    v = ArrayPartition([1.0, 2.0], [3.0])
+    @test v[:] == [1.0, 2.0, 3.0]
+    @test v[:] == vec(v)
+    @test v[:] == collect(v)
+    @test v[1:3] == [1.0, 2.0, 3.0]
+
+    m = ArrayPartition([1.0 2.0], [3.0 4.0 5.0])
+    @test m[:] == [1.0, 2.0, 3.0, 4.0, 5.0]
+    @test m[:] == vec(m)
+    @test m[:] == collect(m)
+    @test m[1:4] == [1.0, 2.0, 3.0, 4.0]
+
+    nested = ArrayPartition(ArrayPartition([1.0], [2.0]), ArrayPartition([3.0, 4.0]))
+    @test nested[:] == [1.0, 2.0, 3.0, 4.0]
+    @test nested[:] == vec(nested)
+    @test nested[:] == collect(nested)
+
+    # (colon, j...) from #239 must still work
+    @test x[:, 1] == (1, 3.0)
+    @test x[:, 2] == (2, 4.0)
+end
+
 # similar partitions
 @inferred similar(x)
 @test similar(x, (4,)) isa ArrayPartition{Float64}