Merge pull request SciML#520 from ChrisRackauckas-Claude/fix-issue-496-gpu-zeromatrix

Fix zeromatrix to preserve GPU array types for ArrayPartition

Author: ChrisRackauckas

src/array_partition.jl

@@ -548,7 +548,11 @@ end
 ## Linear Algebra
 
 function ArrayInterface.zeromatrix(A::ArrayPartition)
-    x = reduce(vcat, vec.(A.x))
+    # Use foldl with explicit init to preserve array type (important for GPU arrays)
+    # Starting with vec of first element ensures the result type matches the input
+    vecs = vec.(A.x)
+    rest = Base.tail(vecs)
+    x = isempty(rest) ? vecs[1] : foldl(vcat, rest; init = vecs[1])
     return x .* x' .* false
 end
 

src/named_array_partition.jl

@@ -174,7 +174,10 @@ end
 #Overwrite ArrayInterface zeromatrix to work with NamedArrayPartitions & implicit solvers within OrdinaryDiffEq
 function ArrayInterface.zeromatrix(A::NamedArrayPartition)
     B = ArrayPartition(A)
-    x = reduce(vcat, vec.(B.x))
+    # Use foldl with explicit init to preserve array type (important for GPU arrays)
+    vecs = vec.(B.x)
+    rest = Base.tail(vecs)
+    x = isempty(rest) ? vecs[1] : foldl(vcat, rest; init = vecs[1])
     return x .* x' .* false
 end