test for broadcasted sparse

Author: CarloLucibello

Project.toml

@@ -4,7 +4,6 @@ version = "1.28.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

@@ -50,16 +50,12 @@ function rrule(::typeof(findnz), v::AbstractSparseVector)
     return (I, V), findnz_pullback
 end
 
-function rrule(::typeof(broadcast), T::Type{<:Number}, x::AbstractSparseArray)
+function rrule(::typeof(Broadcast.broadcasted), T::Type{<:Number}, x::AbstractSparseArray)
     proj = ProjectTo(x)
-    function broadcast_cast_sparse(Δ)
+    function broadcasted_cast_sparse(Δ)
         return NoTangent(), NoTangent(), proj(Δ)         
     end
-    T.(x), broadcast_cast_sparse
+    T.(x), broadcasted_cast_sparse
 end
 
-# These rules help with testing, and won't hurt:
-# They are correct as we always `collect` the primal result as we need that
-# information for the reverse pass
-# ChainRules.rrule(::typeof(SparseMatrixCSC∘broadcast), T::Type{<:Number}, x::AbstractSparseArray) = rrule(broadcast, T, x)
 

test/rulesets/SparseArrays/sparsematrix.jl

@@ -27,3 +27,8 @@ end
     test_rrule(findnz, A ⊢ dA, output_tangent=(zeros(length(I)), zeros(length(J)), V̄))
 end
 
+@testset "broadcasted cast SparseMatrixCSC" begin
+    A = sprand(5, 5, 0.5)
+    test_rrule(_broadcast, Float32, A, rtol=1e-5)
+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