sparsevec test

JuliaDiff · Feb 18, 2022 · 5d5741c · 5d5741c
1 parent 3d98d9d
commit 5d5741c
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/src/rulesets/SparseArrays/sparsematrix.jl b/src/rulesets/SparseArrays/sparsematrix.jl
@@ -44,7 +44,7 @@ function rrule(::typeof(findnz), v::AbstractSparseVector)
     function findnz_pullback(Δ)
         _, V̄ = unthunk(Δ)
         V̄ isa AbstractZero && return (NoTangent(), V̄)
-        return NoTangent(), sparse(I, V̄, n)
+        return NoTangent(), sparsevec(I, V̄, n)
     end
 
     return (I, V), findnz_pullback

diff --git a/test/rulesets/SparseArrays/sparsematrix.jl b/test/rulesets/SparseArrays/sparsematrix.jl
@@ -25,4 +25,11 @@ end
     I, J, V = findnz(A)
     V̄ = rand!(similar(V))
     test_rrule(findnz, A ⊢ dA, output_tangent=(zeros(length(I)), zeros(length(J)), V̄))
+
+    v = sprand(5, 0.5)
+    dv = similar(v)
+    rand!(dv.nzval)
+    I, V = findnz(v)
+    V̄ = rand!(similar(V))
+    test_rrule(findnz, v ⊢ dv, output_tangent=(zeros(length(I)), V̄))
 end