Fix copyto! overload for arraypartitions of scalars

Fixes
SciML · Jun 6, 2024 · efa84b2 · efa84b2
1 parent 80a2a24
commit efa84b2
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/src/array_partition.jl b/src/array_partition.jl
@@ -181,7 +181,11 @@ for type in [AbstractArray, SparseArrays.AbstractCompressedVector, PermutedDimsA
         @assert length(dest) == length(A)
         cur = 1
         @inbounds for i in 1:length(A.x)
-            dest[cur:(cur + length(A.x[i]) - 1)] .= vec(A.x[i])
+            if A.x[i] isa Number
+                dest[cur:(cur + length(A.x[i]) - 1)] .= A.x[i]
+            else
+                dest[cur:(cur + length(A.x[i]) - 1)] .= vec(A.x[i])
+            end
             cur += length(A.x[i])
         end
         dest

diff --git a/test/partitions_test.jl b/test/partitions_test.jl
@@ -266,7 +266,7 @@ begin
 end
 
 @testset "Copy and zero with type changing array" begin
-    # Motivating use case for this is ArrayPartitions of Arrow arrays which are mmap:ed and change type when copied 
+    # Motivating use case for this is ArrayPartitions of Arrow arrays which are mmap:ed and change type when copied
     struct TypeChangingArray{T, N} <: AbstractArray{T, N} end
     Base.copy(::TypeChangingArray{T, N}) where {T, N} = Array{T, N}(undef,
         ntuple(_ -> 0, N))
@@ -281,3 +281,9 @@ end
 @testset "Cartesian indexing" begin
     @test ArrayPartition([1, 2], [3])[1:3, 1] == [1, 2, 3]
 end
+
+@testset "Scalar copyto!" begin
+    u = [2.0,1.0]
+    copyto!(u, ArrayPartition(1.0,-1.2))
+    @test u == [1.0,-1.2]
+end