Type-generic zeromatrix for arraypartition

One part of SciML/OrdinaryDiffEq.jl#2703

Author: ChrisRackauckas

src/array_partition.jl

@@ -430,7 +430,10 @@ end
 
 ## Linear Algebra
 
-ArrayInterface.zeromatrix(A::ArrayPartition) = ArrayInterface.zeromatrix(Vector(A))
+function ArrayInterface.zeromatrix(A::ArrayPartition)
+    x = reduce(vcat,vec.(A.x))
+    x .* x' .* false
+end
 
 function __get_subtypes_in_module(
         mod, supertype; include_supertype = true, all = false, except = [])

test/gpu/arraypartition_gpu.jl

@@ -1,4 +1,4 @@
-using RecursiveArrayTools, CUDA, Test
+using RecursiveArrayTools, ArrayInterface, CUDA, Test
 CUDA.allowscalar(false)
 
 # Test indexing with colon
@@ -21,3 +21,6 @@ fill!(pA, false)
 a = ArrayPartition(([1.0f0] |> cu, [2.0f0] |> cu, [3.0f0] |> cu))
 b = ArrayPartition(([0.0f0] |> cu, [0.0f0] |> cu, [0.0f0] |> cu))
 @. a + b
+
+@test ArrayInterface.zeromatrix(ArrayPartition((CUDA.zeros(2),CUDA.zeros(2)))) isa CuMatrix
+@test size(ArrayInterface.zeromatrix(ArrayPartition((CUDA.zeros(2),CUDA.zeros(2))))) == (4,4)