Add rules for det and logdet of Cholesky

Author: devmotion

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRules"
 uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
-version = "1.29.0"
+version = "1.30.0"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/rulesets/LinearAlgebra/dense.jl

@@ -118,15 +118,15 @@ end
 ##### `det`
 #####
 
-function frule((_, Δx), ::typeof(det), x::AbstractMatrix)
+function frule((_, Δx), ::typeof(det), x::StridedMatrix{<:Number})
     Ω = det(x)
     # TODO Performance optimization: probably there is an efficent
     # way to compute this trace without during the full compution within
     return Ω, Ω * tr(x \ Δx)
 end
 frule((_, Δx), ::typeof(det), x::Number) = (det(x), Δx)
 
-function rrule(::typeof(det), x::Union{Number, AbstractMatrix})
+function rrule(::typeof(det), x::Union{Number, StridedMatrix{<:Number}})
     Ω = det(x)
     function det_pullback(ΔΩ)
         ∂x = x isa Number ? ΔΩ : inv(x)' * dot(Ω, ΔΩ)

src/rulesets/LinearAlgebra/factorization.jl

@@ -551,3 +551,24 @@ function rrule(::typeof(getproperty), F::T, x::Symbol) where {T <: Cholesky}
     end
     return getproperty(F, x), getproperty_cholesky_pullback
 end
+
+# `det` and `logdet` for `Cholesky`
+function rrule(::typeof(det), C::Cholesky)
+    y = det(C)
+    s = conj!((2 * y) ./ _diag_view(C.factors))
+    function det_Cholesky_pullback(ȳ)
+        ΔC = Tangent{typeof(C)}(; factors=Diagonal(ȳ .* s))
+        return NoTangent(), ΔC
+    end
+    return y, det_Cholesky_pullback
+end
+
+function rrule(::typeof(logdet), C::Cholesky)
+    y = logdet(C)
+    s = conj!((2 * one(eltype(C))) ./ _diag_view(C.factors))
+    function logdet_Cholesky_pullback(ȳ)
+        ΔC = Tangent{typeof(C)}(; factors=Diagonal(ȳ .* s))
+        return NoTangent(), ΔC
+    end
+    return y, logdet_Cholesky_pullback
+end

test/rulesets/LinearAlgebra/factorization.jl

@@ -432,5 +432,24 @@ end
             ΔX_symmetric = chol_back_sym(Δ)[2]
             @test sym_back(ΔX_symmetric)[2] ≈ dX_pullback(Δ)[2]
         end
+
+        @testset "det and logdet (uplo=$p)" for p in ['U', 'L']
+            @testset "$op" for op in (det, logdet)
+                @testset "$T" for T in (Float64, ComplexF64)
+                    n = 5
+                    # rand (not randn) so det will be postive, so logdet will be defined
+                    A = 3 * rand(T, (n, n))
+                    X = Cholesky((p === 'U' ? UpperTriangular : LowerTriangular)(A * A' + I))
+                    X̄_acc = Tangent{typeof(X)}(; factors=Diagonal(randn(T, n))) # sensitivity is always a diagonal
+                    test_rrule(op, X ⊢ X̄_acc)
+
+                    # return type
+                    _, op_pullback = rrule(op, X)
+                    X̄ = op_pullback(2.7)[2]
+                    @test X̄ isa Tangent{<:Cholesky}
+                    @test X̄.factors isa Diagonal
+                end
+            end
+        end
     end
 end