Update ext/MatrixAlgebraKitEnzymeExt/MatrixAlgebraKitEnzymeExt.jl

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

Author: kshyatt

ext/MatrixAlgebraKitEnzymeExt/MatrixAlgebraKitEnzymeExt.jl

@@ -372,18 +372,15 @@ for (f!, f_full!, pb!) in (
             # so we do not need to cache it. This may change if future pullbacks
             # depend directly on A!
             nD, V = $f_full!(A.val, alg.val)
-            if TD == Nothing || isa(D, Const)
-                ret = diagview(nD)
-                dret = zero(ret)
-            else
-                ret = D.val
-                copy!(ret, diagview(nD))
-                dret = D.dval
-            end
+            ret = TD == Nothing ? diagview(nD) : copy!(D.val, diagview(nD))
             cache_D = (D.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_D, dret, V))
+            dret = if EnzymeRules.needs_shadow(config)
+                TD == Nothing || isa(D, Const) ? zero(ret) : D.dval
+            else
+                nothing
+            end
+            return EnzymeRules.AugmentedReturn(primal, dret, (cache_D, dret, V))
         end
         function EnzymeRules.reverse(
                 config::EnzymeRules.RevConfigWidth{1},