Add rules for division by Cholesky (#631)

* Add rules for division by Cholesky

* Add tests for new rules

* Increment minor version number

* Increment minor version number

Author: sethaxen

Project.toml

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

src/rulesets/LinearAlgebra/factorization.jl

@@ -578,3 +578,38 @@ function _x_divide_conj_y(x, y)
     z = x / conj(y)
     return iszero(x) ? zero(z) : z
 end
+
+# these rules exists because the primals mutates using `ldiv!` and `rdiv!`
+function rrule(::typeof(\), A::Cholesky, B::AbstractVecOrMat{<:Union{Real,Complex}})
+    U, getproperty_back = rrule(getproperty, A, :U)
+    Z = U' \ B
+    Y = U \ Z
+    project_B = ProjectTo(B)
+    function ldiv_Cholesky_AbsVecOrMat_pullback(ΔY)
+        ∂Z = U' \ ΔY
+        ∂B = U \ ∂Z
+        ∂A = Thunk() do
+            _, Ā = getproperty_back(-add!!(∂Z * Y', Z * ∂B'))
+            return Ā
+        end
+        return NoTangent(), ∂A, project_B(∂B)
+    end
+    return Y, ldiv_Cholesky_AbsVecOrMat_pullback
+end
+
+function rrule(::typeof(/), B::AbstractMatrix{<:Union{Real,Complex}}, A::Cholesky)
+    U, getproperty_back = rrule(getproperty, A, :U)
+    Z = B / U
+    Y = Z / U'
+    project_B = ProjectTo(B)
+    function rdiv_AbstractMatrix_Cholesky_pullback(ΔY)
+        ∂Z = ΔY / U
+        ∂B = ∂Z / U'
+        ∂A = Thunk() do
+            _, Ā = getproperty_back(-add!!(∂Z' * Y, Z' * ∂B))
+            return Ā
+        end
+        return NoTangent(), project_B(∂B), ∂A
+    end
+    return Y, rdiv_AbstractMatrix_Cholesky_pullback
+end

test/rulesets/LinearAlgebra/factorization.jl

@@ -521,5 +521,29 @@ end
                 @test ΔX.factors isa Diagonal && all(iszero, ΔX.factors)
             end
         end
+
+        @testset "\\(::Cholesky, ::AbstractVecOrMat)" begin
+            n = 10
+            for T in (Float64, ComplexF64), sz in (n, (n, 5))
+                A = generate_well_conditioned_matrix(T, n)
+                C = cholesky(A)
+                B = randn(T, sz)
+                # because the rule calls the rrule for getproperty, its rrule is not
+                # completely type-inferrable
+                test_rrule(\, C, B; check_inferred=false)
+            end
+        end
+
+        @testset "/(::AbstractMatrix, ::Cholesky)" begin
+            n = 10
+            for T in (Float64, ComplexF64)
+                A = generate_well_conditioned_matrix(T, n)
+                C = cholesky(A)
+                B = randn(T, 5, n)
+                # because the rule calls the rrule for getproperty, its rrule is not
+                # completely type-inferrable
+                test_rrule(/, B, C; check_inferred=false)
+            end
+        end
     end
 end