A bit of Enzyme cleanup

Author: kshyatt

ext/MatrixAlgebraKitEnzymeExt/MatrixAlgebraKitEnzymeExt.jl

@@ -104,17 +104,13 @@ for (f, pb) in (
             A_is_arg = A_is_arg1 || A_is_arg2
             argval = something(cache_arg, arg.val)
             if !isa(A, Const)
-                if A_is_arg
-                    ΔA = make_zero(A.dval)
-                    $pb(ΔA, Aval, argval, darg)
-                    A.dval .= ΔA
-                else
-                    $pb(A.dval, Aval, argval, darg)
-                end
+                Δ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.dval === arg.dval[1] || make_zero!(arg.dval[1])
-                A.dval === arg.dval[2] || make_zero!(arg.dval[2])
+                A_is_arg1 || make_zero!(arg.dval[1])
+                A_is_arg2 || make_zero!(arg.dval[2])
             end
             return (nothing, nothing, nothing)
         end
@@ -343,7 +339,7 @@ for (f, trunc_f, full_f, pb) in (
             $pb(A.dval, Aval, DVval, dDVtrunc, ind)
         end
         if !isa(DV, Const)
-            if !(A.dval === DV.dval[1])
+            if A.dval !== DV.dval[1]
                 make_zero!(DV.dval)
             else
                 make_zero!(DV.dval[2])
@@ -397,17 +393,11 @@ for (f!, f_full!, pb!) in (
             Aval = nothing
             A_is_arg = !isa(A, Const) && TA <: Diagonal && diagview(A.dval) === D.dval
             if !isa(A, Const)
-                if A_is_arg
-                    ΔA = make_zero(A.dval)
-                    $pb!(ΔA, Aval, (Diagonal(Dval), V), dD)
-                    A.dval .= ΔA
-                else
-                    $pb!(A.dval, Aval, (Diagonal(Dval), V), dD)
-                end
-            end
-            if !isa(D, Const) && !A_is_arg
-                make_zero!(D.dval)
+                Δ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) && !A_is_arg && make_zero!(D.dval)
             return (nothing, nothing, nothing)
         end
     end
@@ -452,17 +442,11 @@ function EnzymeRules.reverse(
     Sval = something(cache_S, S.val)
     A_is_arg = !isa(A, Const) && TA <: Diagonal && diagview(A.dval) === S.dval
     if !isa(A, Const)
-        if A_is_arg
-            ΔA = make_zero(A.dval)
-            svd_vals_pullback!(ΔA, Aval, (U, Diagonal(Sval), Vᴴ), dS)
-            A.dval .= ΔA
-        else
-            svd_vals_pullback!(A.dval, Aval, (U, Diagonal(Sval), Vᴴ), dS)
-        end
-    end
-    if !isa(S, Const) && !A_is_arg
-        make_zero!(S.dval)
+        Δ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) && !A_is_arg && make_zero!(S.dval)
     return (nothing, nothing, nothing)
 end