rrule for broadcasted cast of sparse matrix

Author: CarloLucibello

Project.toml

@@ -4,6 +4,7 @@ version = "1.27.0"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 IrrationalConstants = "92d709cd-6900-40b7-9082-c6be49f344b6"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/rulesets/SparseArrays/sparsematrix.jl

@@ -48,9 +48,6 @@ function rrule(::typeof(findnz), v::AbstractSparseVector)
     n = length(v)
 
     function findnz_pullback(Δ)
-        Δ === NoTangent() && return (NoTangent(), Δ)
-        Δ === ZeroTangent() && return (NoTangent(), Δ)
-        
         _, V̄ = unthunk(Δ)
 
         V̄ === NoTangent() && return (NoTangent(), V̄)
@@ -61,3 +58,11 @@ function rrule(::typeof(findnz), v::AbstractSparseVector)
 
     return (I, V), findnz_pullback
 end
+
+function rrule(::typeof(Broadcast.broadcasted), T::Type{<:Number}, x::AbstractSparseArray)
+    proj = ProjectTo(x)
+    function broadcasted_cast_sparse(Δ)
+        return NoTangent(), NoTangent(), proj(Δ)         
+    end
+    T.(x), broadcasted_cast_sparse
+end

test/rulesets/SparseArrays/sparsematrix.jl

@@ -19,9 +19,17 @@ end
 end
 
 @testset "findnz" begin
-    A = sprand(5, 5, 0.2)
-    test_rrule(findnz, A)
+    A = sprand(5, 5, 0.5)
+    dA = similar(A)
+    rand!(dA.nzval)
+    I, J, V = findnz(A)
+    V̄ = rand!(similar(V))
+    test_rrule(findnz, A ⊢ dA, output_tangent=(NoTangent(), NoTangent(), V̄))
 
-    v = sprand(10, 0.5)
-    test_rrule(findnz, v)
+    # A = sprand(5, 5, 0.5)
+    # test_rrule(findnz, A)
+
+    # v = sprand(10, 0.5)
+    # test_rrule(findnz, v)
 end
+