rrule for broadcasted cast of sparse matrix

Co-authored-by: Miha Zgubic <[email protected]>

Author: CarloLucibello

Project.toml

@@ -16,7 +16,7 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 ChainRulesCore = "1.12"
 ChainRulesTestUtils = "1.5"
 Compat = "3.35"
-FiniteDifferences = "0.12.20"
+FiniteDifferences = "0.12.23"
 IrrationalConstants = "0.1.1"
 JuliaInterpreter = "0.8" # latest is "0.9.1"
 RealDot = "0.1"

src/rulesets/SparseArrays/sparsematrix.jl

@@ -49,3 +49,13 @@ 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
+
+    return T.(x), broadcasted_cast_sparse
+end

test/rulesets/SparseArrays/sparsematrix.jl

@@ -33,3 +33,8 @@ end
     V̄ = rand!(similar(V))
     test_rrule(findnz, v ⊢ dv, output_tangent=(zeros(length(I)), V̄))
 end
+
+@testset "broadcasted cast SparseMatrixCSC" begin
+    A = sprand(5, 5, 0.5)
+    test_rrule(_broadcast, Float32, A, rtol=1e-4)
+end

test/test_helpers.jl

@@ -75,6 +75,28 @@ function ChainRulesCore.rrule(::typeof(make_two_vec), x)
     return make_two_vec(x), make_two_vec_pullback
 end
 
+"""
+Helper function for testing rrule(::typeof(Broadcast.broadcasted), ...)
+
+# Examples
+
+```julia
+using SparseArrays, ChainRulesTestUtils
+
+A = sprand(5, 5, 0.5)
+test_rrule(_broadcast, Float32, A, rtol=1e-5)
+```
+"""
+_broadcast(f::F, x...) where F = broadcast(f, x...)
+
+# we need this due to https://github.com/JuliaDiff/ChainRulesTestUtils.jl/issues/234
+function ChainRulesCore.rrule(::typeof(_broadcast), f::F, args...) where F
+    rr = rrule(Broadcast.broadcasted, f, args...)
+    rr === nothing && return nothing
+    y, pb = rr
+    Broadcast.materialize(Broadcast.instantiate(y)), pb
+end
+
 @testset "test_helpers.jl" begin
 
     @testset "Multiplier" begin