rrule for findnz

Author: CarloLucibello

src/rulesets/SparseArrays/sparsematrix.jl

@@ -10,16 +10,57 @@ function rrule(::typeof(sparse), I::AbstractVector, J::AbstractVector, V::Abstra
     return sparse(I, J, V, m, n, combine), sparse_pullback
 end
 
-function rrule(::Type{T}, A::AbstractMatrix) where T <: SparseMatrixCSC
+function rrule(::Type{T}, A::AbstractMatrix) where T <: AbstractSparseMatrix
     function sparse_pullback(Ω̄)
         return NoTangent(), Ω̄
     end
     return T(A), sparse_pullback
 end
 
-function rrule(::Type{T}, v::AbstractVector) where T <: SparseVector
+function rrule(::Type{T}, v::AbstractVector) where T <: AbstractSparseVector
     function sparse_pullback(Ω̄)
         return NoTangent(), Ω̄
     end
     return T(v), sparse_pullback
 end
+
+function rrule(::typeof(findnz), A::AbstractSparseMatrix)
+    I, J, V = findnz(A)
+    @show I, J, V
+    m, n = size(A)
+
+    function findnz_pullback(Δ)
+        Δ === NoTangent() && return (NoTangent(), Δ)
+        Δ === ZeroTangent() && return (NoTangent(), Δ)
+        
+        _, _, V̄ = unthunk(Δ)
+        
+        V̄ === NoTangent() && return (NoTangent(), V̄)
+        V̄ === ZeroTangent() && return (NoTangent(), V̄)
+        
+        @show I, J, V̄ 
+
+        return NoTangent(), sparse(I, J, V̄, m, n)
+    end
+
+    return (I, J, V), findnz_pullback
+end
+
+function rrule(::typeof(findnz), v::AbstractSparseVector)
+    I, V = findnz(v)
+    n = length(v)
+
+    function findnz_pullback(Δ)
+        Δ === NoTangent() && return (NoTangent(), Δ)
+        Δ === ZeroTangent() && return (NoTangent(), Δ)
+        
+        _, V̄ = unthunk(Δ)
+
+        V̄ === NoTangent() && return (NoTangent(), V̄)
+        V̄ === ZeroTangent() && return (NoTangent(), V̄)
+        
+        return NoTangent(), sparse(I, V̄, n)
+    end
+
+    return (I, V), findnz_pullback
+end

test/rulesets/SparseArrays/sparsematrix.jl

@@ -17,3 +17,11 @@ end
     test_rrule(SparseVector, v)
     test_rrule(SparseVector{Float32}, Float32.(v), rtol=1e-5)
 end
+
+@testset "findnz" begin
+    A = sprand(5, 5, 0.2)
+    test_rrule(findnz, A)
+
+    v = sprand(10, 0.5)
+    test_rrule(findnz, v)
+end