From c32d19b3c39b5ea37a057dde68694d01ab2b1860 Mon Sep 17 00:00:00 2001
From: AshtonSBradley <ashton.bradley@gmail.com>
Date: Thu, 7 May 2026 19:59:29 +1200
Subject: [PATCH] Fix reverse custom rule rooting for sparse wrappers

---
 src/rules/customrules.jl                      |  6 ++-
 .../internal_rules/sparsearrays_rules.jl      | 42 +++++++++++++++++++
 2 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/src/rules/customrules.jl b/src/rules/customrules.jl
index 70a648947f..5c08a64d27 100644
--- a/src/rules/customrules.jl
+++ b/src/rules/customrules.jl
@@ -456,7 +456,11 @@ function enzyme_custom_setup_args(
         root_ty = nothing
         roots_val = if roots_op !== nothing
             root_ty = convert(LLVMType, AnyArray(roots))
-            new_from_original(gutils, roots_op)
+            roots_val = new_from_original(gutils, roots_op)
+            if reverse && B !== nothing
+                roots_val = lookup_value(gutils, roots_val, B)
+            end
+            roots_val
         end
 
 
diff --git a/test/rules/internal_rules/sparsearrays_rules.jl b/test/rules/internal_rules/sparsearrays_rules.jl
index a840896741..35a207aacc 100644
--- a/test/rules/internal_rules/sparsearrays_rules.jl
+++ b/test/rules/internal_rules/sparsearrays_rules.jl
@@ -46,6 +46,31 @@ function test_sparse(M, v, α, β)
     return test_reverse(LinearAlgebra.mul!, Const, (C, Const), (M, Const), (v, Const), (α, Active), (β, Active))
 end
 
+struct SparseWrapperScale{T}
+    λ::T
+end
+
+Base.convert(::Type{Number}, x::SparseWrapperScale) = x.λ
+Base.iszero(x::SparseWrapperScale) = iszero(x.λ)
+
+struct SparseWrapperMatrix{A}
+    A::A
+end
+
+struct SparseWrapperScaledMatrix
+    λ::SparseWrapperScale{Float64}
+    L::SparseWrapperMatrix{SparseMatrixCSC{Float64, Int}}
+end
+
+function sparse_wrapper_mul!(w, L::SparseWrapperScaledMatrix, v)
+    iszero(L.λ) && return lmul!(false, w)
+    α = convert(Number, L.λ)
+    return mul!(w, L.L.A, v, α, false)
+end
+
+const sparse_wrapper_matrix = sparse(Float64[0.0 1.0; 0.0 0.0])
+sparse_wrapper_nonconst_m = SparseWrapperMatrix(sparse_wrapper_matrix)
+
 @testset "SparseArrays spmatvec reverse rule" begin
     Ts = ComplexF64
 
@@ -82,6 +107,23 @@ end
     end
 end
 
+@testset "SparseArrays nested wrapper reverse rule" begin
+    function f(p)
+        u = [3.0, 4.0]
+        du = similar(u)
+        L = SparseWrapperScaledMatrix(SparseWrapperScale(-p[1]), sparse_wrapper_nonconst_m)
+        sparse_wrapper_mul!(du, L, u)
+        return sum(du)
+    end
+
+    p = [1.0]
+    dp = Enzyme.make_zero(p)
+
+    @test f(p) == -4.0
+    Enzyme.autodiff(Enzyme.set_runtime_activity(Enzyme.Reverse), f, Active, Duplicated(p, dp))
+    @test dp ≈ [-4.0]
+end
+
 @testset "SparseArrays spmatmat reverse rule" begin
     Ts = ComplexF64