Avoid undef error

Signed-off-by: ErikQQY <[email protected]>
SciML · Sep 11, 2024 · 8ff79f5 · 8ff79f5
1 parent 3047711
commit 8ff79f5
Showing 1 changed file with 15 additions and 1 deletion.
diff --git a/src/vector_of_array.jl b/src/vector_of_array.jl
@@ -736,7 +736,21 @@ end
 end
 
 @inline function Base.similar(VA::VectorOfArray, dims::N) where {N}
-    VectorOfArray(similar(VA.u, dims))
+    l = length(VA)
+    if dims <= l
+        VectorOfArray([similar(VA[:, i]) for i in 1:l])
+    else
+        VectorOfArray([[similar(VA[:, i]) for i in 1:l]; [similar(VA.u[end]) for _ in (l+1):dims]])
+    end
+end
+
+@inline function Base.similar(VA::VectorOfArray, ::Type{T}, dims::N) where {N, T}
+    l = length(VA)
+    if dims <= l
+        VectorOfArray([similar(VA[:, i], T) for i in 1:l])
+    else
+        VectorOfArray([[similar(VA[:, i], T) for i in 1:l]; [similar(VA.u[end], T) for _ in (l+1):dims]])
+    end
 end
 
 @inline function Base.similar(VA::VectorOfArray{T, N, AT},