Merge branch 'main' of https://github.com/JuliaDiff/ChainRules.jl into spdiag

ElOceanografo · ElOceanografo · commit a2005d1fdb3f · 2023-09-22T14:54:52.000-07:00
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
@@ -21,7 +21,7 @@ jobs:
           - x86
           - x64
     steps:
-      - uses: actions/checkout@v3.5.2
+      - uses: actions/checkout@v4.0.0
       - uses: julia-actions/setup-julia@v1
         with:
           version: ${{ matrix.version }}
diff --git a/.github/workflows/IntegrationTest.yml b/.github/workflows/IntegrationTest.yml
@@ -25,14 +25,14 @@ jobs:
         #    package: {user: JuliaDiff, repo: Diffractor.jl}
 
     steps:
-      - uses: actions/checkout@v3.5.2
+      - uses: actions/checkout@v4.0.0
       - uses: julia-actions/setup-julia@v1
         with:
           version: ${{ matrix.julia-version }}
           arch: x64
       - uses: julia-actions/julia-buildpkg@latest
       - name: Clone Downstream
-        uses: actions/checkout@v3.5.2
+        uses: actions/checkout@v4.0.0
         with:
           repository: ${{ matrix.package.user }}/${{ matrix.package.repo }}
           path: downstream
diff --git a/.github/workflows/JuliaNightly.yml b/.github/workflows/JuliaNightly.yml
@@ -23,7 +23,7 @@ jobs:
           - x86
           - x64
     steps:
-      - uses: actions/checkout@v3.5.2
+      - uses: actions/checkout@v4.0.0
       - uses: julia-actions/setup-julia@v1
         with:
           version: ${{ matrix.version }}
diff --git a/.github/workflows/VersionVigilante_pull_request.yml b/.github/workflows/VersionVigilante_pull_request.yml
@@ -6,7 +6,7 @@ jobs:
   VersionVigilante:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v3.5.2
+      - uses: actions/checkout@v4.0.0
       - uses: julia-actions/setup-julia@latest
       - name: VersionVigilante.main
         id: versionvigilante_main
diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "ChainRules"
 uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
-version = "1.49.0"
+version = "1.54.0"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -14,8 +14,10 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 RealDot = "c1ae055f-0cd5-4b69-90a6-9a35b1a98df9"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+SparseInverseSubset = "dc90abb0-5640-4711-901d-7e5b23a2fada"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
+SuiteSparse = "4607b0f0-06f3-5cda-b6b1-a6196a1729e9"
 
 [compat]
 Adapt = "3.4.0"
@@ -28,6 +30,7 @@ IrrationalConstants = "0.1.1, 0.2"
 JLArrays = "0.1"
 JuliaInterpreter = "0.8,0.9"
 RealDot = "0.1"
+SparseInverseSubset = "0.1"
 StaticArrays = "1.2"
 StructArrays = "0.6.11"
 julia = "1.6"
diff --git a/src/rulesets/Base/arraymath.jl b/src/rulesets/Base/arraymath.jl
@@ -210,7 +210,13 @@ end # VERSION
 ##### `muladd`
 #####
 
-function frule((_, ΔA, ΔB, Δz), ::typeof(muladd), A, B, z)
+function frule(
+    (_, ΔA, ΔB, Δz),
+    ::typeof(muladd),
+    A::AbstractVecOrMat{<:CommutativeMulNumber},
+    B::AbstractVecOrMat{<:CommutativeMulNumber},
+    z::Union{CommutativeMulNumber, AbstractVecOrMat{<:CommutativeMulNumber}}
+)
     Ω = muladd(A, B, z)
     return Ω, ΔA * B .+ A * ΔB .+ Δz
 end
@@ -342,20 +348,43 @@ function rrule(::typeof(\), A::AbstractVecOrMat{<:Real}, B::AbstractVecOrMat{<:R
     project_B = ProjectTo(B)
 
     Y = A \ B
+
     function backslash_pullback(ȳ)
         Ȳ = unthunk(ȳ)
+        
+        Ȳf = Ȳ
+        @static if VERSION >= v"1.9"
+            # Need to ensure Ȳ is an array since since https://github.com/JuliaLang/julia/pull/44358
+            if !isa(Ȳ, AbstractArray)
+                Ȳf = [Ȳ]
+            end
+        end
+        Yf = Y
+        @static if VERSION >= v"1.9" 
+            # Need to ensure Yf is an array since since https://github.com/JuliaLang/julia/pull/44358
+            if !isa(Y, AbstractArray)
+                Yf = [Y]
+            end
+        end
+        #@info "vars" typeof(Ȳ) typeof(Y) typeof(Yf) typeof(A) typeof(B)
         ∂A = @thunk begin
-            B̄ = A' \ Ȳ
+            B̄ = A' \ Ȳf
             Ā = -B̄ * Y'
-            Ā = add!!(Ā, (B - A * Y) * B̄' / A')
-            Ā = add!!(Ā, A' \ Y * (Ȳ' - B̄'A))
+            t = (B - A * Y) * B̄'
+            @static if VERSION >= v"1.9" 
+                # Need to ensure t is an array since since https://github.com/JuliaLang/julia/pull/44358
+                if !isa(t, AbstractArray)
+                    t = [t]
+                end
+            end
+            Ā = add!!(Ā, t / A')
+            Ā = add!!(Ā, A' \ Yf * (Ȳ' - B̄'A))
             project_A(Ā)
         end
-        ∂B = @thunk project_B(A' \ Ȳ)
+        ∂B = @thunk project_B(A' \ Ȳf)
         return NoTangent(), ∂A, ∂B
     end
     return Y, backslash_pullback
-
 end
 
 #####
diff --git a/src/rulesets/Base/base.jl b/src/rulesets/Base/base.jl
@@ -3,8 +3,8 @@
 
 @scalar_rule copysign(y, x) (ifelse(signbit(x)!=signbit(y), -one(y), +one(y)), NoTangent())
 
-@scalar_rule one(x) zero(x)
-@scalar_rule zero(x) zero(x)
+@scalar_rule one(x) ZeroTangent()
+@scalar_rule zero(x) ZeroTangent()
 @scalar_rule transpose(x) true
 
 # `adjoint`
diff --git a/src/rulesets/Base/fastmath_able.jl b/src/rulesets/Base/fastmath_able.jl
@@ -169,8 +169,8 @@ let
         @scalar_rule x / y (one(x) / y, -(Ω / y))
         
         ## many-arg +
-        function frule((_, Δx, Δy...), ::typeof(+), x::Number, ys::Number...)
-            +(x, ys...), +(Δx, Δy...)
+        function frule(Δs, ::typeof(+), x::Number, ys::Number...)
+            +(x, ys...), +(Base.tail(Δs)...)
         end
         
         function rrule(::typeof(+), x::Number, ys::Number...)
diff --git a/src/rulesets/Base/indexing.jl b/src/rulesets/Base/indexing.jl
@@ -1,14 +1,6 @@
-#####
-##### getindex(::Tuple)
-#####
-
-function frule((_, ẋ), ::typeof(getindex), x::Tuple, i::Integer)
-    return x[i], ẋ[i]
-end
-
-function frule((_, ẋ), ::typeof(getindex), x::Tuple, i)
-    y = x[i]
-    return y, Tangent{typeof(y)}(ẋ[i]...)
+# Int rather than Int64/Integer is intentional
+function frule((_, ẋ), ::typeof(getfield), x::Tuple, i::Int)
+    return x.i, ẋ.i
 end
 
 "for a given tuple type, returns a Val{N} where N is the length of the tuple"
@@ -77,20 +69,52 @@ end
 """
     ∇getindex(x, dy, inds...)
 
-For the `rrule` of `y = x[inds...]`, this function is roughly 
+For the `rrule` of `y = x[inds...]`, this function is roughly
 `setindex(zero(x), dy, inds...)`, returning the array `dx`.
 Differentiable. Includes `ProjectTo(x)(dx)`.
 """
-function ∇getindex(x::AbstractArray, dy, inds...)
+function ∇getindex(x::AbstractArray{T,N}, dy, inds...) where {T,N}
     # `to_indices` removes any logical indexing, colons, CartesianIndex etc,
     # leaving just Int / AbstractVector of Int
     plain_inds = Base.to_indices(x, inds)
-    dx = _setindex_zero(x, dy, plain_inds...)
-    ∇getindex!(dx, dy, plain_inds...)
+    dx = if plain_inds isa NTuple{N, Int} && T<:Number
+        # scalar indexing
+        OneElement(dy, plain_inds, axes(x))
+    else  # some from slicing (potentially noncontigous)
+        dx = _setindex_zero(x, dy, plain_inds...)
+        ∇getindex!(dx, dy, plain_inds...)
+    end
     return ProjectTo(x)(dx)  # since we have x, may as well do this inside, not in rules
 end
 ∇getindex(x::AbstractArray, z::AbstractZero, inds...) = z
 
+"""
+    OneElement(val, ind, axes) <: AbstractArray
+
+Extremely simple `struct` used for the gradient of scalar `getindex`.
+"""
+struct OneElement{T,N,I,A} <: AbstractArray{T,N}
+  val::T
+  ind::I
+  axes::A
+  OneElement(val::T, ind::I, axes::A) where {T<:Number, I<:NTuple{N,Int}, A<:NTuple{N,AbstractUnitRange}} where {N} = new{T,N,I,A}(val, ind, axes)
+end
+Base.size(A::OneElement) = map(length, A.axes)
+Base.axes(A::OneElement) = A.axes
+Base.getindex(A::OneElement{T,N}, i::Vararg{Int,N}) where {T,N} = ifelse(i==A.ind, A.val, zero(T))
+
+function ChainRulesCore.add!!(xs::AbstractArray{<:Any,N}, oe::OneElement{<:Any,N}) where {N}
+    if !ChainRulesCore.is_inplaceable_destination(xs)
+        xs = collect(xs)
+    end
+    xs[oe.ind...] += oe.val
+    return xs
+end
+
+Base.:(+)(xs::AbstractArray, oe::OneElement) = add!!(copy(xs), oe)
+Base.:(+)(oe::OneElement, xs::AbstractArray) = +(xs, oe)
+Base.:(+)(oe1::OneElement, oe2::OneElement) = +(collect(oe1), oe2)
+
 """
     _setindex_zero(x, dy, inds...)
 
@@ -159,29 +183,6 @@ function ∇getindex!(dx::AbstractGPUArray, dy, inds...)
     return dx
 end
 
-#####
-##### first, tail
-#####
-
-function frule((_, ẋ), ::typeof(first), x::Tuple)
-    return first(x), first(ẋ)
-end
-
-function rrule(::typeof(first), x::T) where {T<:Tuple}
-    first_back(dy) = (NoTangent(), Tangent{T}(ntuple(j -> j == 1 ? dy : NoTangent(), _tuple_N(T))...))
-    return first(x), first_back
-end
-
-function frule((_, ẋ), ::typeof(Base.tail), x::Tuple)
-    y = Base.tail(x)
-    return y, Tangent{typeof(y)}(Base.tail(ẋ)...)
-end
-
-function rrule(::typeof(Base.tail), x::T) where {T<:Tuple}
-    tail_pullback(dy) = (NoTangent(), Tangent{T}(NoTangent(), dy...))
-    return Base.tail(x), tail_pullback
-end
-
 #####
 ##### view
 #####
diff --git a/src/rulesets/Base/nondiff.jl b/src/rulesets/Base/nondiff.jl
@@ -441,6 +441,8 @@ end
 
 @non_differentiable unescape_string(::AbstractString)
 @non_differentiable unescape_string(::IO, ::AbstractString)
+@non_differentiable unique(::AbstractArray{<:Union{Regex, AbstractChar, AbstractString}})
+@non_differentiable unique(::AbstractArray{Symbol})
 @non_differentiable unmark(::IO)
 @non_differentiable unsafe_string(::Cstring)
 @non_differentiable uppercase(::AbstractString)
diff --git a/src/rulesets/LinearAlgebra/structured.jl b/src/rulesets/LinearAlgebra/structured.jl
@@ -1,5 +1,6 @@
 # Structured matrices
 using LinearAlgebra: AbstractTriangular
+using SparseInverseSubset
 
 # Matrix wrapper types that we know are square and are thus potentially invertible. For
 # these we can use simpler definitions for `/` and `\`.
diff --git a/src/rulesets/SparseArrays/sparsematrix.jl b/src/rulesets/SparseArrays/sparsematrix.jl
@@ -50,32 +50,90 @@ function rrule(::typeof(findnz), v::AbstractSparseVector)
     return (I, V), findnz_pullback
 end
 
-function _spdiagm_back(p, ȳ)
-    k, v = p
-    d = diag(unthunk(ȳ), k)[1:length(v)] # handle if diagonal was smaller than matrix
-    return Tangent{typeof(p)}(second = d)
+if VERSION < v"1.7"
+    #=
+    The method below for `logabsdet(F::UmfpackLU)` is required to calculate the (log) 
+    determinants of sparse matrices, but was not defined prior to Julia v1.7. In order
+    for the rrules for the determinants of sparse matrices below to work, they need to be
+    able to compute the primals as well, so this import from the future is included. For
+    more recent versions of Julia, this definition lives in:
+    julia/stdlib/SuiteSparse/src/umfpack.jl
+    =#
+    using SuiteSparse.UMFPACK: UmfpackLU
+
+    # compute the sign/parity of a permutation
+    function _signperm(p)
+        n = length(p)
+        result = 0
+        todo = trues(n)
+        while any(todo)
+            k = findfirst(todo)
+            todo[k] = false
+            result += 1 # increment element count
+            j = p[k]
+            while j != k
+                result += 1 # increment element count
+                todo[j] = false
+                j = p[j]
+            end
+            result += 1 # increment cycle count
+        end
+        return ifelse(isodd(result), -1, 1)
+    end
+
+    function LinearAlgebra.logabsdet(F::UmfpackLU{T, TI}) where {T<:Union{Float64,ComplexF64},TI<:Union{Int32, Int64}} 
+        n = checksquare(F)
+        issuccess(F) || return log(zero(real(T))), zero(T)
+        U = F.U
+        Rs = F.Rs
+        p = F.p
+        q = F.q
+        s = _signperm(p)*_signperm(q)*one(real(T))
+        P = one(T)
+        abs_det = zero(real(T))
+        @inbounds for i in 1:n
+            dg_ii = U[i, i] / Rs[i]
+            P *= sign(dg_ii)
+            abs_det += log(abs(dg_ii))
+        end
+        return abs_det, s * P
+    end
 end
 
-function rrule(::typeof(spdiagm), m::Integer, n::Integer, kv::Pair{<:Integer,<:AbstractVector}...)
-    function diagm_pullback(Δ)
-        _, ȳ = unthunk(Δ)
-        return (NoTangent(), NoTangent(), NoTangent(), _spdiagm_back.(kv, Ref(ȳ))...)
+
+function rrule(::typeof(logabsdet), x::SparseMatrixCSC)
+    F = cholesky(x)
+    L, D, U, P = SparseInverseSubset.get_ldup(F)
+    Ω = logabsdet(D)
+    function logabsdet_pullback(ΔΩ)
+        (Δy, Δsigny) = ΔΩ
+        (_, signy) = Ω
+        f = signy' * Δsigny
+        imagf = f - real(f)
+        g = real(Δy) + imagf
+        Z, P = sparseinv(F, depermute=true)
+        ∂x = g * Z'
+        return (NoTangent(), ∂x)
     end
-    return spdiagm(m, n, kv...), diagm_pullback
+    return Ω, logabsdet_pullback
 end
 
-function rrule(::typeof(spdiagm), kv::Pair{<:Integer,<:AbstractVector}...)
-    function diagm_pullback(Δ)
-        _, ȳ = unthunk(Δ)
-        return (NoTangent(), _spdiagm_back.(kv, Ref(ȳ))...)
+function rrule(::typeof(logdet), x::SparseMatrixCSC)
+    Ω = logdet(x)
+    function logdet_pullback(ΔΩ)
+        Z, p = sparseinv(x, depermute=true)
+        ∂x = ΔΩ * Z'
+        return (NoTangent(), ∂x)
     end
-    return spdiagm(kv...), diagm_pullback
+    return Ω, logdet_pullback
 end
 
-function rrule(::typeof(spdiagm), v::AbstractVector)
-    function diagm_pullback(Δ)
-        _, ȳ = unthunk(Δ)
-        return (NoTangent(), diag(ȳ))
+function rrule(::typeof(det), x::SparseMatrixCSC)
+    Ω = det(x)
+    function det_pullback(ΔΩ)
+        Z, _ = sparseinv(x, depermute=true)
+        ∂x = Z' * dot(Ω, ΔΩ)
+        return (NoTangent(), ∂x)
     end
-    return spdiagm(v), diagm_pullback
+    return Ω, det_pullback
 end
diff --git a/test/rulesets/Base/array.jl b/test/rulesets/Base/array.jl
diff --git a/test/rulesets/Base/arraymath.jl b/test/rulesets/Base/arraymath.jl
diff --git a/test/rulesets/Base/indexing.jl b/test/rulesets/Base/indexing.jl
diff --git a/test/rulesets/SparseArrays/sparsematrix.jl b/test/rulesets/SparseArrays/sparsematrix.jl
