sparse: Mutating arrays is not supported while SparseMatrixCSC works

[zhehaoli1999]

Hi there, I am trying to differentiate through sparse matrix:

using LinearAlgebra
using SparseArrays
using ChainRulesCore
using Zygote

Zygote.@adjoint function SparseMatrixCSC{T,N}(arr) where {T,N}
    SparseMatrixCSC{T,N}(arr), Δ -> (collect(Δ),)
  end

function test1(a)
    A = sparse([1, 1, 2, 2],[1, 2, 1, 2],a, 2, 2)
    return sum(A)
end 

a = [1.0,2.0,3.0,4.0]
gradient(test1, a)

Then I got error:

ERROR: LoadError: Mutating arrays is not supported
Stacktrace:
 [1] error(::String) at ./error.jl:33
 [2] (::Zygote.var"#399#400")(::Nothing) at /root/.julia/packages/Zygote/CgsVi/src/lib/array.jl:58
 [3] (::Zygote.var"#2265#back#401"{Zygote.var"#399#400"})(::Nothing) at /root/.julia/packages/ZygoteRules/OjfTt/src/adjoint.jl:59
 [4] sparse! at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.5/SparseArrays/src/sparsematrix.jl:862 [inlined]
 [5] (::typeof(∂(sparse!)))(::FillArrays.Fill{Float64,2,Tuple{Base.OneTo{Int64},Base.OneTo{Int64}}}) at /root/.julia/packages/Zygote/CgsVi/src/compiler/interface2.jl:0
 [6] sparse at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.5/SparseArrays/src/sparsematrix.jl:703 [inlined]
 [7] (::typeof(∂(sparse)))(::FillArrays.Fill{Float64,2,Tuple{Base.OneTo{Int64},Base.OneTo{Int64}}}) at /root/.julia/packages/Zygote/CgsVi/src/compiler/interface2.jl:0
 [8] sparse at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.5/SparseArrays/src/sparsematrix.jl:892 [inlined]
 [9] (::typeof(∂(sparse)))(::FillArrays.Fill{Float64,2,Tuple{Base.OneTo{Int64},Base.OneTo{Int64}}}) at /root/.julia/packages/Zygote/CgsVi/src/compiler/interface2.jl:0
 [10] test1 at /root/codes/test_zygote/test_sparse.jl:11 [inlined]
 [11] (::typeof(∂(test1)))(::Float64) at /root/.julia/packages/Zygote/CgsVi/src/compiler/interface2.jl:0
 [12] (::Zygote.var"#41#42"{typeof(∂(test1))})(::Float64) at /root/.julia/packages/Zygote/CgsVi/src/compiler/interface.jl:41
 [13] gradient(::Function, ::Array{Float64,1}) at /root/.julia/packages/Zygote/CgsVi/src/compiler/interface.jl:59
 [14] top-level scope at /root/codes/test_zygote/test_sparse.jl:16
in expression starting at /root/codes/test_zygote/test_sparse.jl:16

Then I changed sparse to SparseMatrixCSC， it works:

using LinearAlgebra
using SparseArrays
using ChainRulesCore
using Zygote

Zygote.@adjoint function SparseMatrixCSC{T,N}(arr) where {T,N}
    SparseMatrixCSC{T,N}(arr), Δ -> (collect(Δ),)
  end

function test2(a)
    A = SparseMatrixCSC(2,2,[1, 3, 5], [1, 2, 1, 2],a)
    return sum(A)
end 

a = [1.0,2.0,3.0,4.0]
gradient(test2, a)

output

([1.0, 1.0, 1.0, 1.0],)

As far as I am concerned, calling sparse will also create SparseMatrixCSC, so

1. why SparseMatrixCSC works while sparse doesn't?
2. What is the root cause of "Mutating arrays is not supported" when calling sparse?
3. How can I directly differentiate sparse?

Thx for any reply.

[zhehaoli1999]

I think I found the answer to the first and the second question above.
The reason is that sparse will add the values (which is mutating the array) if two indices are the same in the input index array, while SparseMatrixCSC not.
See the examples below:

julia> sparse([1,1,1,2,2,2],[1,1,2,2,2,2], [1,2,3,4,5,6], 2,2)
2×2 SparseMatrixCSC{Int64,Int64} with 3 stored entries:
  [1, 1]  =  3
  [1, 2]  =  3
  [2, 2]  =  15

julia> SparseMatrixCSC(2,2,[1, 3, 7], [1, 1, 1, 2,2, 2],[1, 2, 3, 4, 5, 6])
2×2 SparseMatrixCSC{Int64,Int64} with 6 stored entries:Error showing value of type SparseMatrixCSC{Int64,Int64}:
ERROR: BoundsError: attempt to access 4-element Array{Int64,1} at index [1:6]

Now there is only one question left: How can I directly differentiate sparse?

[DhairyaLGandhi]

Ref #762

Could you try with that?

[zhehaoli1999]

Hi, thx for replying.
After reading through this pull request, I tried: (Which is just copying from the commits of the pull request)

using LinearAlgebra
using SparseArrays
using ChainRulesCore
using Zygote:@adjoint, @ignore, @nograd, gradient

#####################################
# 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
##################################################

function test1(a)
    A = sparse([1, 1, 2, 2],[1, 2, 1, 2],a, 2, 2)
    return sum(A)
end 

a = [1.0,2.0,3.0,4.0]
gradient(test1, a)

The result is still:

ERROR: LoadError: Mutating arrays is not supported
Stacktrace:
 [1] error(::String) at ./error.jl:33
 [2] (::Zygote.var"#372#373")(::Nothing) at /root/.julia/packages/Zygote/KpME9/src/lib/array.jl:58
 [3] (::Zygote.var"#2261#back#374"{Zygote.var"#372#373"})(::Nothing) at /root/.julia/packages/ZygoteRules/OjfTt/src/adjoint.jl:59
 [4] sparse! at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.5/SparseArrays/src/sparsematrix.jl:862 [inlined]
 [5] (::typeof(∂(sparse!)))(::FillArrays.Fill{Float64,2,Tuple{Base.OneTo{Int64},Base.OneTo{Int64}}}) at /root/.julia/packages/Zygote/KpME9/src/compiler/interface2.jl:0
 [6] sparse at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.5/SparseArrays/src/sparsematrix.jl:703 [inlined]
 [7] (::typeof(∂(sparse)))(::FillArrays.Fill{Float64,2,Tuple{Base.OneTo{Int64},Base.OneTo{Int64}}}) at /root/.julia/packages/Zygote/KpME9/src/compiler/interface2.jl:0
 [8] sparse at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.5/SparseArrays/src/sparsematrix.jl:892 [inlined]
 [9] (::typeof(∂(sparse)))(::FillArrays.Fill{Float64,2,Tuple{Base.OneTo{Int64},Base.OneTo{Int64}}}) at /root/.julia/packages/Zygote/KpME9/src/compiler/interface2.jl:0
 [10] test1 at /root/codes/test_zygote/test_sparse.jl:39 [inlined]
 [11] (::typeof(∂(test1)))(::Float64) at /root/.julia/packages/Zygote/KpME9/src/compiler/interface2.jl:0
 [12] (::Zygote.var"#41#42"{typeof(∂(test1))})(::Float64) at /root/.julia/packages/Zygote/KpME9/src/compiler/interface.jl:40
 [13] gradient(::Function, ::Array{Float64,1}) at /root/.julia/packages/Zygote/KpME9/src/compiler/interface.jl:49
 [14] top-level scope at /root/codes/test_zygote/test_sparse.jl:49
in expression starting at /root/codes/test_zygote/test_sparse.jl:49

I may misunderstand your words. Am I doing the wrong way? Thanks. @DhairyaLGandhi