Update ext/MatrixAlgebraKitEnzymeExt/MatrixAlgebraKitEnzymeExt.jl

Co-authored-by: Jutho <Jutho@users.noreply.github.com>

Author: kshyatt

ext/MatrixAlgebraKitEnzymeExt/MatrixAlgebraKitEnzymeExt.jl

@@ -430,7 +430,15 @@ function EnzymeRules.augmented_primal(
     cache_S = (S.val !== ret) || EnzymeRules.overwritten(config)[3] ? copy(ret) : nothing
     primal = EnzymeRules.needs_primal(config) ? ret : nothing
     shadow = EnzymeRules.needs_shadow(config) ? dret : nothing
-    return EnzymeRules.AugmentedReturn(primal, shadow, (cache_S, dret, U, Vᴴ))
+    ret = TS == Nothing ? diagview(nS) : copy!(S.val, diagview(nS))
+    cache_S = (S.val !== ret) || EnzymeRules.overwritten(config)[3] ? copy(ret) : nothing
+    primal = EnzymeRules.needs_primal(config) ? ret : nothing
+    dret = if EnzymeRules.needs_shadow(config)
+        TS == Nothing || isa(S, Const) : zero(ret) : S.dval
+    else
+        nothing
+    end
+    return EnzymeRules.AugmentedReturn(primal, dret, (cache_S, dret, U, Vᴴ))
 end
 function EnzymeRules.reverse(
         config::EnzymeRules.RevConfigWidth{1},