diff --git a/src/internal_rules/linalg.jl b/src/internal_rules/linalg.jl index 616189bd46..1702bef935 100644 --- a/src/internal_rules/linalg.jl +++ b/src/internal_rules/linalg.jl @@ -295,6 +295,79 @@ function EnzymeRules.forward( end end +function EnzymeRules.augmented_primal( + config::EnzymeRules.RevConfig, + func::Const{typeof(ldiv!)}, + ::Type{RT}, + fact::Annotation{<:Cholesky}, + B::Annotation{<:AbstractVecOrMat}; + kwargs..., + ) where {RT} + if B isa Const + retval = func.val(fact.val, B.val; kwargs...) + primal = EnzymeRules.needs_primal(config) ? retval : nothing + return EnzymeRules.AugmentedReturn(primal, nothing, nothing) + end + + cache_fact = if EnzymeRules.overwritten(config)[2] + Cholesky(copy(fact.val.factors), fact.val.uplo, fact.val.info) + else + fact.val + end + L = cache_fact.L + U = cache_fact.U + + ldiv!(L, B.val) + cache_Lout = fact isa Const ? nothing : copy(B.val) + ldiv!(U, B.val) + cache_Uout = fact isa Const ? nothing : copy(B.val) + + primal = EnzymeRules.needs_primal(config) ? B.val : nothing + shadow = EnzymeRules.needs_shadow(config) ? B.dval : nothing + + return EnzymeRules.AugmentedReturn(primal, shadow, (cache_fact, cache_Lout, cache_Uout)) +end + +function EnzymeRules.reverse( + config::EnzymeRules.RevConfig, + func::Const{typeof(ldiv!)}, + ::Type{RT}, + cache, + fact::Annotation{<:Cholesky}, + B::Annotation{<:AbstractVecOrMat}; + kwargs..., + ) where {RT} + if B isa Const + return (nothing, nothing) + end + + cache_fact, cache_Lout, cache_Uout = cache + L = cache_fact.L + U = cache_fact.U + N = EnzymeRules.width(config) + + ntuple(Val(N)) do i + Base.@_inline_meta + dB = N == 1 ? B.dval : B.dval[i] + zU = adjoint(U) \ dB + zL = adjoint(L) \ zU + if !(fact isa Const) + dfact = N == 1 ? fact.dval : fact.dval[i] + dU = zU * adjoint(cache_Uout) + dL = zL * adjoint(cache_Lout) + if cache_fact.uplo == 'U' + dfact.factors .-= triu!(dU + adjoint(dL)) + else + dfact.factors .-= tril!(dL + adjoint(dU)) + end + end + dB .= zL + return nothing + end + + return (nothing, nothing) +end + _zero_unused_elements!(X, ::UpperTriangular) = triu!(X) _zero_unused_elements!(X, ::LowerTriangular) = tril!(X) @@ -503,4 +576,3 @@ end # partial derivative of the determinant is the matrix of cofactors EnzymeRules.@easy_rule(LinearAlgebra.det(A::AbstractMatrix), (cofactor(A),)) - diff --git a/test/rules/internal_rules/linear_algebra_rules.jl b/test/rules/internal_rules/linear_algebra_rules.jl index d438f8b476..7867171b25 100644 --- a/test/rules/internal_rules/linear_algebra_rules.jl +++ b/test/rules/internal_rules/linear_algebra_rules.jl @@ -523,6 +523,41 @@ end end end +function chol_ldiv_weighted(fact, B, W)::Float64 + B = copy(B) + ldiv!(fact, B) + return LinearAlgebra.dot(W, B) +end + +@testset "Reverse Cholesky ldiv!" begin + F = Float64[ + 1.4 0.2 + 0.0 1.5 + ] + b = Float64[0.7, -0.2] + w = Float64[0.3, -0.4] + fact = LinearAlgebra.Cholesky(copy(F), 'U', 0) + test_reverse(chol_ldiv_weighted, Active, (fact, Duplicated), (b, Duplicated), (w, Const)) + test_reverse(chol_ldiv_weighted, Active, (fact, BatchDuplicated), (b, BatchDuplicated), (w, Const)) + test_reverse(chol_ldiv_weighted, Active, (fact, Const), (b, Duplicated), (w, Const)) + test_reverse(chol_ldiv_weighted, Active, (fact, Duplicated), (b, Const), (w, Const)) + + F = Float64[ + 1.4 0.0 + 0.2 1.5 + ] + B = Float64[ + 0.7 -0.2 + 0.1 0.5 + ] + W = Float64[ + 0.3 -0.4 + 0.2 0.6 + ] + fact = LinearAlgebra.Cholesky(copy(F), 'L', 0) + test_reverse(chol_ldiv_weighted, Active, (fact, Duplicated), (B, Duplicated), (W, Const)) +end + function chol_upper(x) x = reshape(x, 4, 4) x = parent(cholesky(Hermitian(x)).U)