From 81ff85ba4ec03cd8d1a2a9ef35a63121c676db9a Mon Sep 17 00:00:00 2001
From: Jesse Perla <jesseperla@gmail.com>
Date: Fri, 15 May 2026 18:58:30 -0700
Subject: [PATCH 1/2] feat: add Cholesky ldiv reverse rule

---
 src/internal_rules/linalg.jl                  | 72 ++++++++++++++++++-
 .../internal_rules/linear_algebra_rules.jl    | 35 +++++++++
 2 files changed, 106 insertions(+), 1 deletion(-)

diff --git a/src/internal_rules/linalg.jl b/src/internal_rules/linalg.jl
index 616189bd46..f6a79b0fb0 100644
--- a/src/internal_rules/linalg.jl
+++ b/src/internal_rules/linalg.jl
@@ -295,6 +295,77 @@ 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)
+
+    for i in 1:N
+        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
+    end
+
+    return (nothing, nothing)
+end
+
 
 _zero_unused_elements!(X, ::UpperTriangular) = triu!(X)
 _zero_unused_elements!(X, ::LowerTriangular) = tril!(X)
@@ -503,4 +574,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)

From 827ad0bbfc0f65d2502ef89f4eff8ee8593c7601 Mon Sep 17 00:00:00 2001
From: Jesse Perla <jesseperla@gmail.com>
Date: Sat, 16 May 2026 09:47:20 -0700
Subject: [PATCH 2/2] refactor: use ntuple for batched reverse loop

Replace `for i in 1:N` with `ntuple(Val(N)) do i ... end` in the
Cholesky ldiv! reverse rule, matching the forward rule pattern.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
---
 src/internal_rules/linalg.jl | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/internal_rules/linalg.jl b/src/internal_rules/linalg.jl
index f6a79b0fb0..1702bef935 100644
--- a/src/internal_rules/linalg.jl
+++ b/src/internal_rules/linalg.jl
@@ -346,7 +346,8 @@ function EnzymeRules.reverse(
     U = cache_fact.U
     N = EnzymeRules.width(config)
 
-    for i in 1:N
+    ntuple(Val(N)) do i
+        Base.@_inline_meta
         dB = N == 1 ? B.dval : B.dval[i]
         zU = adjoint(U) \ dB
         zL = adjoint(L) \ zU
@@ -361,6 +362,7 @@ function EnzymeRules.reverse(
             end
         end
         dB .= zL
+        return nothing
     end
 
     return (nothing, nothing)