Small fixes

Author: Katharine Hyatt

ext/MatrixAlgebraKitEnzymeExt/MatrixAlgebraKitEnzymeExt.jl

@@ -379,6 +379,7 @@ for (f!, f_full!, pb!) in (
             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
+            # on 1.10, Enzyme can get confused about whether it needs the shadow
             dret = if EnzymeRules.needs_shadow(config)
                 TD == Nothing || isa(D, Const) ? zero(ret) : D.dval
             else
@@ -397,13 +398,16 @@ for (f!, f_full!, pb!) in (
             ) where {RT}
             cache_D, dD, V = cache
             Dval = something(cache_D, D.val)
+            # on 1.10, Enzyme can get confused about whether it needs the shadow
+            # replace the dret with the arg.dval in the case it's nothing
+            dDval = something(dD, D.dval)
             # A is  NOT used in the pullback, so we assign Aval = nothing
             # to trigger an error in case the pullback is modified to directly
             # use A (so that whoever does this is forced to handle caching A
             # appropriately here)
             Aval = nothing
             if !isa(A, Const)
-                $pb!(A.dval, Aval, (Diagonal(Dval), V), dD)
+                $pb!(A.dval, Aval, (Diagonal(Dval), V), dDval)
             end
             !isa(D, Const) && make_zero!(D.dval)
             return (nothing, nothing, nothing)
@@ -426,8 +430,9 @@ function EnzymeRules.augmented_primal(
     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
+    # on 1.10, Enzyme can get confused about whether it needs the shadow
     dret = if EnzymeRules.needs_shadow(config)
-        TS == Nothing || isa(S, Const):zero(ret):S.dval
+        TS == Nothing || isa(S, Const) ? zero(ret) : S.dval
     else
         nothing
     end
@@ -449,8 +454,11 @@ function EnzymeRules.reverse(
     # appropriately here)
     Aval = nothing
     Sval = something(cache_S, S.val)
+    # on 1.10, Enzyme can get confused about whether it needs the shadow
+    # replace the dret with the arg.dval in the case it's nothing
+    dSval = something(dS, S.dval)
     if !isa(A, Const)
-        svd_vals_pullback!(A.dval, Aval, (U, Diagonal(Sval), Vᴴ), dS)
+        svd_vals_pullback!(A.dval, Aval, (U, Diagonal(Sval), Vᴴ), dSval)
     end
     !isa(S, Const) && make_zero!(S.dval)
     return (nothing, nothing, nothing)