diff --git a/ext/MatrixAlgebraKitEnzymeExt/MatrixAlgebraKitEnzymeExt.jl b/ext/MatrixAlgebraKitEnzymeExt/MatrixAlgebraKitEnzymeExt.jl
index 85072c949..ba6512b5b 100644
--- a/ext/MatrixAlgebraKitEnzymeExt/MatrixAlgebraKitEnzymeExt.jl
+++ b/ext/MatrixAlgebraKitEnzymeExt/MatrixAlgebraKitEnzymeExt.jl
@@ -43,37 +43,6 @@ using LinearAlgebra
 # other provided input variable.
 #--------------------------------------------
 
-# this rule is necessary for now as without it,
-# a segfault occurs both on 1.10 and 1.12 -- likely
-# a deeper internal bug
-function EnzymeRules.augmented_primal(
-        config::EnzymeRules.RevConfigWidth{1},
-        func::Const{typeof(copy_input)},
-        ::Type{RT},
-        f::Annotation,
-        A::Annotation
-    ) where {RT}
-    ret = func.val(f.val, A.val)
-    primal = EnzymeRules.needs_primal(config) ? ret : nothing
-    shadow = EnzymeRules.needs_shadow(config) ? zero(A.dval) : nothing
-    return EnzymeRules.AugmentedReturn(primal, shadow, shadow)
-end
-
-function EnzymeRules.reverse(
-        config::EnzymeRules.RevConfigWidth{1},
-        func::Const{typeof(copy_input)},
-        ::Type{RT},
-        cache,
-        f::Annotation,
-        A::Annotation
-    ) where {RT}
-    copy_shadow = cache
-    if !isa(A, Const) && !isnothing(copy_shadow)
-        A.dval .+= copy_shadow
-    end
-    return (nothing, nothing)
-end
-
 # two-argument factorizations like LQ, QR, EIG
 for (f, pb) in (
         (qr_full!, qr_pullback!),
@@ -101,9 +70,14 @@ for (f, pb) in (
             # if arg.val == ret, the annotation must be Duplicated or DuplicatedNoNeed
             # if arg isa Const, ret may still be modified further down the call graph so we should
             # copy it to protect ourselves
-            cache_arg = (arg.val !== ret) || EnzymeRules.overwritten(config)[3] ? copy.(ret) : nothing
-            dret = if EnzymeRules.needs_shadow(config)
-                (TA == Nothing && TB == Nothing) || isa(arg, Const) ? zero.(ret) : arg.dval
+            A_is_arg1 = !isa(A, Const) && A.val === arg.val[1]
+            A_is_arg2 = !isa(A, Const) && A.val === arg.val[2]
+            A_is_arg = A_is_arg1 || A_is_arg2
+            cache_arg = (arg.val !== ret && !A_is_arg) || EnzymeRules.overwritten(config)[3] ? copy.(ret) : nothing
+            dret = if EnzymeRules.needs_shadow(config) && ((TA == Nothing && TB == Nothing) || isa(arg, Const))
+                make_zero.(ret)
+            elseif EnzymeRules.needs_shadow(config)
+                arg.dval
             else
                 nothing
             end
@@ -125,11 +99,19 @@ for (f, pb) in (
             # use A (so that whoever does this is forced to handle caching A
             # appropriately here)
             Aval = nothing
+            A_is_arg1 = !isa(A, Const) && A.dval === arg.dval[1]
+            A_is_arg2 = !isa(A, Const) && A.dval === arg.dval[2]
+            A_is_arg = A_is_arg1 || A_is_arg2
             argval = something(cache_arg, arg.val)
             if !isa(A, Const)
-                $pb(A.dval, Aval, argval, darg)
+                ΔA = A_is_arg ? make_zero(A.dval) : A.dval
+                $pb(ΔA, Aval, argval, darg)
+                A_is_arg && (A.dval .= ΔA)
+            end
+            if !isa(arg, Const)
+                A_is_arg1 || make_zero!(arg.dval[1])
+                A_is_arg2 || make_zero!(arg.dval[2])
             end
-            !isa(arg, Const) && make_zero!(arg.dval)
             return (nothing, nothing, nothing)
         end
     end
@@ -356,7 +338,13 @@ for (f, trunc_f, full_f, pb) in (
         if !isa(A, Const)
             $pb(A.dval, Aval, DVval, dDVtrunc, ind)
         end
-        !isa(DV, Const) && make_zero!(DV.dval)
+        if !isa(DV, Const)
+            if A.dval !== DV.dval[1]
+                make_zero!(DV.dval)
+            else
+                make_zero!(DV.dval[2])
+            end
+        end
         return (nothing, nothing, nothing)
     end
 end
@@ -392,10 +380,10 @@ for (f!, f_full!, pb!) in (
                 func::Const{typeof($f!)},
                 ::Type{RT},
                 cache,
-                A::Annotation,
+                A::Annotation{TA},
                 D::Annotation,
                 alg::Const{<:MatrixAlgebraKit.AbstractAlgorithm},
-            ) where {RT}
+            ) where {RT, TA}
             cache_D, dD, V = cache
             Dval = something(cache_D, D.val)
             # A is  NOT used in the pullback, so we assign Aval = nothing
@@ -403,10 +391,13 @@ for (f!, f_full!, pb!) in (
             # use A (so that whoever does this is forced to handle caching A
             # appropriately here)
             Aval = nothing
+            A_is_arg = !isa(A, Const) && TA <: Diagonal && diagview(A.dval) === D.dval
             if !isa(A, Const)
-                $pb!(A.dval, Aval, (Diagonal(Dval), V), dD)
+                ΔA = A_is_arg ? make_zero(A.dval) : A.dval
+                $pb!(ΔA, Aval, (Diagonal(Dval), V), dD)
+                A_is_arg && (A.dval .= ΔA)
             end
-            !isa(D, Const) && make_zero!(D.dval)
+            !isa(D, Const) && !A_is_arg && make_zero!(D.dval)
             return (nothing, nothing, nothing)
         end
     end
@@ -438,10 +429,10 @@ function EnzymeRules.reverse(
         func::Const{typeof(svd_vals!)},
         ::Type{RT},
         cache,
-        A::Annotation,
+        A::Annotation{TA},
         S::Annotation,
         alg::Const{<:MatrixAlgebraKit.AbstractAlgorithm},
-    ) where {RT}
+    ) where {RT, TA}
     cache_S, dS, U, Vᴴ = cache
     # A is  NOT used in the pullback, so we assign Aval = nothing
     # to trigger an error in case the pullback is modified to directly
@@ -449,10 +440,13 @@ function EnzymeRules.reverse(
     # appropriately here)
     Aval = nothing
     Sval = something(cache_S, S.val)
+    A_is_arg = !isa(A, Const) && TA <: Diagonal && diagview(A.dval) === S.dval
     if !isa(A, Const)
-        svd_vals_pullback!(A.dval, Aval, (U, Diagonal(Sval), Vᴴ), dS)
+        ΔA = A_is_arg ? make_zero(A.dval) : A.dval
+        svd_vals_pullback!(ΔA, Aval, (U, Diagonal(Sval), Vᴴ), dS)
+        A_is_arg && (A.dval .= ΔA)
     end
-    !isa(S, Const) && make_zero!(S.dval)
+    !isa(S, Const) && !A_is_arg && make_zero!(S.dval)
     return (nothing, nothing, nothing)
 end
 
diff --git a/ext/MatrixAlgebraKitMooncakeExt/MatrixAlgebraKitMooncakeExt.jl b/ext/MatrixAlgebraKitMooncakeExt/MatrixAlgebraKitMooncakeExt.jl
index ef32d6de4..363356ff4 100644
--- a/ext/MatrixAlgebraKitMooncakeExt/MatrixAlgebraKitMooncakeExt.jl
+++ b/ext/MatrixAlgebraKitMooncakeExt/MatrixAlgebraKitMooncakeExt.jl
@@ -15,19 +15,6 @@ using LinearAlgebra
 
 Mooncake.tangent_type(::Type{<:MatrixAlgebraKit.AbstractAlgorithm}) = Mooncake.NoTangent
 
-@is_primitive Mooncake.DefaultCtx Mooncake.ReverseMode Tuple{typeof(copy_input), Any, Any}
-function Mooncake.rrule!!(::CoDual{typeof(copy_input)}, f_df::CoDual, A_dA::CoDual)
-    Ac = copy_input(Mooncake.primal(f_df), Mooncake.primal(A_dA))
-    Ac_dAc = Mooncake.zero_fcodual(Ac)
-    dAc = Mooncake.tangent(Ac_dAc)
-    function copy_input_pb(::NoRData)
-        Mooncake.increment!!(Mooncake.tangent(A_dA), dAc)
-        return NoRData(), NoRData(), NoRData()
-    end
-    return Ac_dAc, copy_input_pb
-end
-
-Mooncake.@zero_derivative Mooncake.DefaultCtx Tuple{typeof(initialize_output), Any, Any, Any}
 # two-argument in-place factorizations like LQ, QR, EIG
 for (f!, f, pb, adj) in (
         (:qr_full!, :qr_full, :qr_pullback!, :qr_adjoint),
@@ -53,13 +40,33 @@ for (f!, f, pb, adj) in (
             arg2c = copy(arg2)
             $f!(A, args, Mooncake.primal(alg_dalg))
             function $adj(::NoRData)
-                copy!(A, Ac)
-                $pb(dA, A, (arg1, arg2), (darg1, darg2))
-                copy!(arg1, arg1c)
-                copy!(arg2, arg2c)
-                zero!(darg1)
-                zero!(darg2)
-                return NoRData(), NoRData(), NoRData(), NoRData()
+                # DON'T copy Ac to A if A === one
+                # of the output args -- this can
+                # mess up the pullback because
+                # generally the args are used there
+                if !(A === arg1 || A === arg2)
+                    copy!(A, Ac)
+                    $pb(dA, A, (arg1, arg2), (darg1, darg2))
+                else
+                    ΔA = zero(A)
+                    $pb(ΔA, A, (arg1, arg2), (darg1, darg2))
+                    dA .= ΔA
+                end
+                if A === arg1
+                    copy!(A, Ac)
+                    zero!(darg2)
+                    copy!(arg2, arg2c)
+                elseif A === arg2
+                    copy!(A, Ac)
+                    zero!(darg1)
+                    copy!(arg1, arg1c)
+                else
+                    zero!(darg1)
+                    zero!(darg2)
+                    copy!(arg2, arg2c)
+                    copy!(arg1, arg1c)
+                end
+                return ntuple(Returns(NoRData()), 4)
             end
             return args_dargs, $adj
         end
@@ -140,9 +147,19 @@ for (f!, f, f_full, pb, adj) in (
             copy!(D, diagview(DV[1]))
             V = DV[2]
             function $adj(::NoRData)
-                $pb(dA, A, DV, dD)
-                copy!(D, Dc)
-                zero!(dD)
+                if A !== D
+                    $pb(dA, A, DV, dD)
+                else
+                    ΔA = zero(A)
+                    $pb(ΔA, A, DV, dD)
+                    dA .= A
+                end
+                if A !== D
+                    zero!(dD)
+                    copy!(D, Dc)
+                else
+                    copy!(A, Ac)
+                end
                 return NoRData(), NoRData(), NoRData(), NoRData()
             end
             return D_dD, $adj
@@ -199,15 +216,27 @@ for f in (:eig, :eigh)
             # not for nested structs with various fields (like Diagonal{Complex})
             output_codual = Mooncake.zero_fcodual(output)
             function $f_adjoint!(dy::Tuple{NoRData, NoRData, <:Real})
-                copy!(A, Ac)
                 Dtrunc, Vtrunc, ϵ = Mooncake.primal(output_codual)
                 dDtrunc_, dVtrunc_, dϵ = Mooncake.tangent(output_codual)
                 _warn_pullback_truncerror(dy[3])
                 D′, dD′ = arrayify(Dtrunc, dDtrunc_)
                 V′, dV′ = arrayify(Vtrunc, dVtrunc_)
-                $f_trunc_pullback!(dA, A, (D′, V′), (dD′, dV′))
-                copy!(DV[1], DVc[1])
-                copy!(DV[2], DVc[2])
+                D, dD = arrayify(DV[1], dDV[1])
+                V, dV = arrayify(DV[2], dDV[2])
+                copy!(A, Ac)
+                if !(A === D || A === V)
+                    $f_trunc_pullback!(dA, A, (D′, V′), (dD′, dV′))
+                else
+                    ΔA = zero(A)
+                    $f_trunc_pullback!(ΔA, A, (D′, V′), (dD′, dV′))
+                    dA .= ΔA
+                end
+                if A === D
+                    copy!(DV[2], DVc[2])
+                else
+                    copy!(DV[1], DVc[1])
+                    copy!(DV[2], DVc[2])
+                end
                 zero!(dD′)
                 zero!(dV′)
                 return NoRData(), NoRData(), NoRData(), NoRData()
@@ -239,12 +268,22 @@ for f in (:eig, :eigh)
                 _warn_pullback_truncerror(dϵ)
 
                 # compute pullbacks
-                $f_pullback!(dA, Ac, DV, dDVtrunc, ind)
-                zero!.(dDVtrunc) # since this is allocated in this function this is probably not required
-
+                if !(A === DV[1] || A === DV[2])
+                    $f_pullback!(dA, Ac, DV, dDVtrunc, ind)
+                else
+                    ΔA = zero(A)
+                    $f_pullback!(ΔA, Ac, DV, dDVtrunc, ind)
+                    dA .= ΔA
+                end
                 # restore state
                 copy!(A, Ac)
-                copy!.(DV, DVc)
+                if A === DV[1]
+                    copy!(DV[2], DVc[2])
+                    zero!(dDV[2])
+                else
+                    copy!.(DV, DVc)
+                    zero!.(dDV)
+                end
 
                 return ntuple(Returns(NoRData()), 4)
             end
@@ -351,12 +390,23 @@ for f in (:eig, :eigh)
             dDVtrunc = last.(arrayify.(DVtrunc, Mooncake.tangent(DVtrunc_dDVtrunc)))
             function $f_adjoint!(::NoRData)
                 # compute pullbacks
-                $f_pullback!(dA, Ac, DV, dDVtrunc, ind)
-                zero!.(dDV)
+                if !(A === DV[1] || A === DV[2])
+                    $f_pullback!(dA, Ac, DV, dDVtrunc, ind)
+                else
+                    ΔA = zero(A)
+                    $f_pullback!(ΔA, Ac, DV, dDVtrunc, ind)
+                    dA .= ΔA
+                end
 
                 # restore state
                 copy!(A, Ac)
-                copy!.(DV, DVc)
+                if A === DV[1]
+                    copy!(DV[2], DVc[2])
+                    zero!(dDV[2])
+                else
+                    copy!.(DV, DVc)
+                    zero!.(dDV)
+                end
 
                 return ntuple(Returns(NoRData()), 4)
             end
diff --git a/test/enzyme/eig.jl b/test/enzyme/eig.jl
index 898d773a8..949129eac 100644
--- a/test/enzyme/eig.jl
+++ b/test/enzyme/eig.jl
@@ -15,5 +15,7 @@ for T in (BLASFloats..., GenericFloats...)
     TestSuite.seed_rng!(1234)
     if !is_buildkite
         TestSuite.test_enzyme_eig(T, (m, m); atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
+        AT = Diagonal{T, Vector{T}}
+        TestSuite.test_enzyme_eig(AT, (m, m); atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
     end
 end
diff --git a/test/enzyme/eigh.jl b/test/enzyme/eigh.jl
index d32db3dd5..64c796fc6 100644
--- a/test/enzyme/eigh.jl
+++ b/test/enzyme/eigh.jl
@@ -14,6 +14,8 @@ m = 19
 for T in (BLASFloats..., GenericFloats...)
     TestSuite.seed_rng!(1234)
     if !is_buildkite
-        TestSuite.test_enzyme_eigh(T, (m, m); atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
+        #TestSuite.test_enzyme_eigh(T, (m, m); atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
+        AT = Diagonal{T, Vector{T}}
+        TestSuite.test_enzyme_eigh(AT, (m, m); atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
     end
 end
diff --git a/test/enzyme/lq.jl b/test/enzyme/lq.jl
index f7ae2ebf7..7c747529d 100644
--- a/test/enzyme/lq.jl
+++ b/test/enzyme/lq.jl
@@ -15,5 +15,7 @@ for T in (BLASFloats..., GenericFloats...), n in (17, m, 23)
     TestSuite.seed_rng!(1234)
     if !is_buildkite
         TestSuite.test_enzyme_lq(T, (m, n); atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
+        AT = Diagonal{T, Vector{T}}
+        m == n && TestSuite.test_enzyme_lq(AT, (m, m); atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
     end
 end
diff --git a/test/enzyme/orthnull.jl b/test/enzyme/orthnull.jl
index eaeae8400..086873d3f 100644
--- a/test/enzyme/orthnull.jl
+++ b/test/enzyme/orthnull.jl
@@ -15,5 +15,7 @@ for T in (BLASFloats..., GenericFloats...), n in (17, m, 23)
     TestSuite.seed_rng!(1234)
     if !is_buildkite
         TestSuite.test_enzyme_orthnull(T, (m, n); atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
+        AT = Diagonal{T, Vector{T}}
+        m == n && TestSuite.test_enzyme_orthnull(AT, (m, m); atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
     end
 end
diff --git a/test/enzyme/polar.jl b/test/enzyme/polar.jl
index 6ab965ac1..183086adb 100644
--- a/test/enzyme/polar.jl
+++ b/test/enzyme/polar.jl
@@ -15,5 +15,7 @@ for T in (BLASFloats..., GenericFloats...), n in (17, m, 23)
     TestSuite.seed_rng!(1234)
     if !is_buildkite
         TestSuite.test_enzyme_polar(T, (m, n); atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
+        AT = Diagonal{T, Vector{T}}
+        #m == n && TestSuite.test_enzyme_polar(AT, (m, m); atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
     end
 end
diff --git a/test/enzyme/qr.jl b/test/enzyme/qr.jl
index 728e267d3..2d8b9e7e1 100644
--- a/test/enzyme/qr.jl
+++ b/test/enzyme/qr.jl
@@ -15,5 +15,7 @@ for T in (BLASFloats..., GenericFloats...), n in (17, m, 23)
     TestSuite.seed_rng!(1234)
     if !is_buildkite
         TestSuite.test_enzyme_qr(T, (m, n); atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
+        AT = Diagonal{T, Vector{T}}
+        m == n && TestSuite.test_enzyme_qr(AT, (m, m); atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
     end
 end
diff --git a/test/enzyme/svd.jl b/test/enzyme/svd.jl
index 6143f61e4..e4aaa7aa1 100644
--- a/test/enzyme/svd.jl
+++ b/test/enzyme/svd.jl
@@ -15,5 +15,7 @@ for T in (BLASFloats..., GenericFloats...), n in (17, m, 23)
     TestSuite.seed_rng!(1234)
     if !is_buildkite
         TestSuite.test_enzyme_svd(T, (m, n); atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
+        AT = Diagonal{T, Vector{T}}
+        m == n && TestSuite.test_enzyme_svd(AT, (m, m); atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
     end
 end
diff --git a/test/mooncake/eig.jl b/test/mooncake/eig.jl
index b313f9b2f..a0e606941 100644
--- a/test/mooncake/eig.jl
+++ b/test/mooncake/eig.jl
@@ -15,5 +15,7 @@ for T in (BLASFloats..., GenericFloats...)
     TestSuite.seed_rng!(1234)
     if !is_buildkite
         TestSuite.test_mooncake_eig(T, (m, m); atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
+        AT = Diagonal{T, Vector{T}}
+        TestSuite.test_mooncake_eig(AT, (m, m); atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
     end
 end
diff --git a/test/mooncake/eigh.jl b/test/mooncake/eigh.jl
index 800dbaa05..e39f68316 100644
--- a/test/mooncake/eigh.jl
+++ b/test/mooncake/eigh.jl
@@ -15,5 +15,7 @@ for T in (BLASFloats..., GenericFloats...)
     TestSuite.seed_rng!(1234)
     if !is_buildkite
         TestSuite.test_mooncake_eigh(T, (m, m); atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
+        AT = Diagonal{T, Vector{T}}
+        TestSuite.test_mooncake_eigh(AT, m; atol = m * m * TestSuite.precision(T), rtol = m * m * TestSuite.precision(T))
     end
 end
diff --git a/test/mooncake/lq.jl b/test/mooncake/lq.jl
index 0f05f85ab..42d0fdb6b 100644
--- a/test/mooncake/lq.jl
+++ b/test/mooncake/lq.jl
@@ -15,5 +15,9 @@ for T in (BLASFloats..., GenericFloats...), n in (17, m, 23)
     TestSuite.seed_rng!(1234)
     if !is_buildkite
         TestSuite.test_mooncake_lq(T, (m, n); atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
+        if m == n
+            AT = Diagonal{T, Vector{T}}
+            TestSuite.test_mooncake_lq(AT, (m, m); atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
+        end
     end
 end
diff --git a/test/mooncake/orthnull.jl b/test/mooncake/orthnull.jl
index 09e3a28cc..370454b55 100644
--- a/test/mooncake/orthnull.jl
+++ b/test/mooncake/orthnull.jl
@@ -15,5 +15,9 @@ for T in (BLASFloats..., GenericFloats...), n in (17, m, 23)
     TestSuite.seed_rng!(1234)
     if !is_buildkite
         TestSuite.test_mooncake_orthnull(T, (m, n); atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
+        if m == n
+            AT = Diagonal{T, Vector{T}}
+            TestSuite.test_mooncake_orthnull(AT, m; atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
+        end
     end
 end
diff --git a/test/mooncake/polar.jl b/test/mooncake/polar.jl
index 1faf3c104..6442b7b33 100644
--- a/test/mooncake/polar.jl
+++ b/test/mooncake/polar.jl
@@ -17,5 +17,10 @@ for T in (BLASFloats..., GenericFloats...), n in (17, m, 23)
         atol = rtol = m * n * TestSuite.precision(T)
         m >= n && TestSuite.test_mooncake_left_polar(T, (m, n); atol, rtol)
         n >= m && TestSuite.test_mooncake_right_polar(T, (m, n); atol, rtol)
+        #=if m == n
+            AT = Diagonal{T, Vector{T}}
+            TestSuite.test_mooncake_left_polar(AT, m; atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
+            TestSuite.test_mooncake_right_polar(AT, m; atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
+        end=# # broken due to pullback
     end
 end
diff --git a/test/mooncake/qr.jl b/test/mooncake/qr.jl
index 17415e8df..bbb9a8d17 100644
--- a/test/mooncake/qr.jl
+++ b/test/mooncake/qr.jl
@@ -15,5 +15,9 @@ for T in (BLASFloats..., GenericFloats...), n in (17, m, 23)
     TestSuite.seed_rng!(1234)
     if !is_buildkite
         TestSuite.test_mooncake_qr(T, (m, n); atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
+        if m == n
+            AT = Diagonal{T, Vector{T}}
+            TestSuite.test_mooncake_qr(AT, (m, m); atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
+        end
     end
 end
diff --git a/test/mooncake/svd.jl b/test/mooncake/svd.jl
index d2d40df42..f096fdb8e 100644
--- a/test/mooncake/svd.jl
+++ b/test/mooncake/svd.jl
@@ -15,5 +15,9 @@ for T in (BLASFloats..., GenericFloats...), n in (17, m, 23)
     TestSuite.seed_rng!(1234)
     if !is_buildkite
         TestSuite.test_mooncake_svd(T, (m, n); atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
+        if m == n
+            AT = Diagonal{T, Vector{T}}
+            TestSuite.test_mooncake_svd(AT, m; atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
+        end
     end
 end
diff --git a/test/testsuite/TestSuite.jl b/test/testsuite/TestSuite.jl
index b4fb93008..36ae68304 100644
--- a/test/testsuite/TestSuite.jl
+++ b/test/testsuite/TestSuite.jl
@@ -96,6 +96,12 @@ function instantiate_rank_deficient_matrix(T, sz; trunc = truncrank(div(min(sz..
     return mul!(A, V, C)
 end
 
+function instantiate_rank_deficient_matrix(::Type{T}, sz; trunc = truncrank(div(min(sz...), 2))) where {T <: Diagonal}
+    A = instantiate_matrix(eltype(T), sz)
+    V, C = left_orth!(A; trunc)
+    return Diagonal(diag(mul!(A, V, C)))
+end
+
 include("ad_utils.jl")
 
 include("projections.jl")
diff --git a/test/testsuite/enzyme/orthnull.jl b/test/testsuite/enzyme/orthnull.jl
index 842fdc2f4..9fae30e5e 100644
--- a/test/testsuite/enzyme/orthnull.jl
+++ b/test/testsuite/enzyme/orthnull.jl
@@ -37,7 +37,7 @@ function test_enzyme_left_orth(
             test_reverse(call_and_zero!, RT, (left_orth!, Const), (A, TA), (alg, Const); atol, rtol, fdm, output_tangent = ΔVC)
         end
 
-        if m >= n
+        if m >= n && !(T <: Diagonal)
             @testset "polar" begin
                 A = instantiate_matrix(T, sz)
                 alg = MatrixAlgebraKit.select_algorithm(left_orth!, A, :polar)
@@ -71,7 +71,7 @@ function test_enzyme_right_orth(
             test_reverse(call_and_zero!, RT, (right_orth!, Const), (A, TA), (alg, Const); atol, rtol, fdm, output_tangent = ΔCVᴴ)
         end
 
-        if m <= n
+        if m <= n && !(T <: Diagonal)
             @testset "polar" begin
                 A = instantiate_matrix(T, sz)
                 alg = MatrixAlgebraKit.select_algorithm(right_orth!, A, :polar)
diff --git a/test/testsuite/mooncake/eig.jl b/test/testsuite/mooncake/eig.jl
index 88042df28..aad4b4881 100644
--- a/test/testsuite/mooncake/eig.jl
+++ b/test/testsuite/mooncake/eig.jl
@@ -31,6 +31,12 @@ function test_mooncake_eig_full(
             rng, eig_full, A, alg;
             mode = Mooncake.ReverseMode, output_tangent, atol, rtol
         )
+        if T <: Diagonal{<:Complex}
+            Mooncake.TestUtils.test_rule(
+                rng, eig_full!, A, (A, DV[2]), alg;
+                mode = Mooncake.ReverseMode, output_tangent, atol, rtol
+            )
+        end
         Mooncake.TestUtils.test_rule(
             rng, call_and_zero!, eig_full!, A, alg;
             mode = Mooncake.ReverseMode, output_tangent, atol, rtol, is_primitive = false
@@ -57,6 +63,12 @@ function test_mooncake_eig_vals(
             rng, eig_vals, A, alg;
             mode = Mooncake.ReverseMode, output_tangent, atol, rtol
         )
+        if T <: Diagonal{<:Complex}
+            Mooncake.TestUtils.test_rule(
+                rng, eig_vals!, A, A.diag, alg;
+                mode = Mooncake.ReverseMode, output_tangent, atol, rtol
+            )
+        end
         Mooncake.TestUtils.test_rule(
             rng, call_and_zero!, eig_vals!, A, alg;
             mode = Mooncake.ReverseMode, output_tangent, atol, rtol, is_primitive = false
diff --git a/test/testsuite/mooncake/orthnull.jl b/test/testsuite/mooncake/orthnull.jl
index 48e4ad921..e89700868 100644
--- a/test/testsuite/mooncake/orthnull.jl
+++ b/test/testsuite/mooncake/orthnull.jl
@@ -43,7 +43,7 @@ function test_mooncake_left_orth(
             )
         end
 
-        if m >= n
+        if m >= n && !(T <: Diagonal)
             @testset "polar" begin
                 alg = MatrixAlgebraKit.select_algorithm(left_orth!, A, :polar)
                 VC = left_orth(A, alg)
@@ -91,7 +91,7 @@ function test_mooncake_right_orth(
             )
         end
 
-        if m <= n
+        if m <= n && !(T <: Diagonal)
             @testset "polar" begin
                 alg = MatrixAlgebraKit.select_algorithm(right_orth!, A, :polar)
                 CVᴴ = right_orth(A, alg)