Merge pull request #193 from mateuszbaran/mbaran/similar-array-partition-fix

Fix shaped similar for ArrayPartition

Author: ChrisRackauckas

src/array_partition.jl

@@ -19,17 +19,29 @@ end
 
 Base.similar(A::ArrayPartition{T,S}) where {T,S} = ArrayPartition{T,S}(similar.(A.x))
 
-# ignore dims since array partitions are vectors
-Base.similar(A::ArrayPartition, dims::NTuple{N,Int}) where {N} = similar(A)
+# return ArrayPartition when possible, otherwise next best thing of the correct size
+function Base.similar(A::ArrayPartition, dims::NTuple{N,Int}) where {N}
+    if dims == size(A)
+        return similar(A)
+    else
+        return similar(A.x[1], eltype(A), dims)
+    end
+end
 
 # similar array partition of common type
 @inline function Base.similar(A::ArrayPartition, ::Type{T}) where {T}
     N = npartitions(A)
     ArrayPartition(i->similar(A.x[i], T), N)
 end
 
-# ignore dims since array partitions are vectors
-Base.similar(A::ArrayPartition, ::Type{T}, dims::NTuple{N,Int}) where {T,N} = similar(A, T)
+# return ArrayPartition when possible, otherwise next best thing of the correct size
+function Base.similar(A::ArrayPartition, ::Type{T}, dims::NTuple{N,Int}) where {T,N}
+    if dims == size(A)
+        return similar(A, T)
+    else
+        return similar(A.x[1], T, dims)
+    end
+end
 
 # similar array partition with different types
 function Base.similar(A::ArrayPartition, ::Type{T}, ::Type{S}, R::DataType...) where {T, S}

test/partitions_test.jl

@@ -61,9 +61,11 @@ x = ArrayPartition([1, 2], [3.0, 4.0])
 
 # similar partitions
 @inferred similar(x)
-@inferred similar(x, (2, 2))
+@test similar(x, (4,)) isa ArrayPartition{Float64}
+@test (@inferred similar(x, (2, 2))) isa AbstractMatrix{Float64}
 @inferred similar(x, Int)
-@inferred similar(x, Int, (2, 2))
+@test similar(x, Int, (4,)) isa ArrayPartition{Int}
+@test (@inferred similar(x, Int, (2, 2))) isa AbstractMatrix{Int}
 # @inferred similar(x, Int, Float64)
 
 # zero