basic sparse handling

Project.toml

@@ -16,6 +16,7 @@ MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"

src/lib/array.jl

@@ -1,6 +1,7 @@
 using Random, FillArrays, AbstractFFTs
 using FillArrays: AbstractFill, getindex_value
 using Base.Broadcast: broadcasted, broadcast_shape
+using SparseArrays
 using Distributed: pmap, AbstractWorkerPool
 
 @adjoint (::Type{T})(::UndefInitializer, args...) where T<:Array = T(undef, args...), Δ -> nothing
@@ -957,3 +958,32 @@ end
  back(Δ::AbstractArray) = (nothing, getindex.(_back.(Δ), 1))
  return Fill(y, size(r)), back
 end
+
+# Sparse Arrays
+
+@adjoint function SparseMatrixCSC{T,N}(arr) where {T,N}
+ SparseMatrixCSC{T,N}(arr), Δ -> (collect(Δ),)
+end
+
+@adjoint function SparseVector{T,N}(v) where {T,N}
+ SparseVector{T,N}(v), Δ -> (collect(Δ),)
+end
+
+@adjoint diagm(x::AbstractSparseArray) = diagm(x), Δ -> (diag(Δ), )
+
+@adjoint function Broadcast.broadcasted(::Type{Float32}, a::AbstractSparseArray{T,N}) where {T,N}
+ Float32.(a), Δ -> (nothing, T.(Δ), )
+end
+
+@adjoint Matrix(a::AbstractSparseArray) = Matrix(a), Δ -> (Δ,)
+
+@adjoint function SparseArrays.spdiagm(x::Pair...)
+ ks = first.(x)
+ SparseArrays.spdiagm(x...), Δ -> begin
+ tuple((k => diag(Δ, k) for k in ks)...)
+ end
+end
+
+@adjoint Pair(x,y) = Pair(x,y), Δ -> (nothing, Δ.second)
+
+@nograd issymmetric

test/gradcheck.jl

@@ -2,6 +2,8 @@ using Zygote, Test, Random, LinearAlgebra, Statistics, FillArrays,
  AbstractFFTs, FFTW, Distances
 using Zygote: gradient
 using Base.Broadcast: broadcast_shape
+using Distributed: pmap
+using SparseArrays
 using Distributed: pmap, CachingPool, workers
 import FiniteDifferences
 
@@ -1571,6 +1573,12 @@ end
  end
 end
 
+@testset "Sparse" begin
+ @test gradtest(x -> sum(sparse(x)), rand(Float32, 3,3))
+ @test gradtest(x -> sum(sparse(x)), rand(Float32, 3)) # test vectors also
+ @test gradcheck(x -> sum(diagm(x)), sparse(rand(3)))
+end
+
 @testset "broadcasted($op, Array, Bool)" for op in (+,-,*)
  @testset "with $bool and sizes $s" for s in ((4,), (2,3)), bool in (true,false)
  r = rand(Int8, s) .+ 0.0