Update kzeros and kones for SubArray type

JuliaSmoothOptimizers · May 5, 2021 · ac55536 · ac55536
1 parent e4f2835
commit ac55536
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/src/krylov_utils.jl b/src/krylov_utils.jl
@@ -171,14 +171,13 @@ end
 
 Create an AbstractVector of storage type `S` of length `n` only composed of zero.
 """
-@inline kzeros(S, n) = fill!(S(undef, n), zero(eltype(S)))
-
+@inline kzeros(S, n) = S <: SubArray ? fill!(S.types[1](undef, n), zero(eltype(S))) : fill!(S(undef, n), zero(eltype(S)))
 """
     v = kones(S, n)
 
 Create an AbstractVector of storage type `S` of length `n` only composed of one.
 """
-@inline kones(S, n) = fill!(S(undef, n), one(eltype(S)))
+@inline kones(S, n) = S <: SubArray ? fill!(S.types[1](undef, n), one(eltype(S))) : fill!(similar(S, n), one(eltype(S)))
 
 @inline display(iter, verbose) = (verbose > 0) && (mod(iter, verbose) == 0)
 

diff --git a/test/test_aux.jl b/test/test_aux.jl
@@ -95,4 +95,10 @@
   # test kzeros and kones
   Krylov.kzeros(Vector{Float64}, 10) == zeros(10)
   Krylov.kones(Vector{Float64}, 10) == ones(10)
+
+  a = rand(10)
+  b = view(a, 1:4)
+  S = typeof(b)
+  Krylov.kzeros(S, 10) == zeros(10)
+  Krylov.kones(S, 10) == ones(10)
 end