fix: zero out shadows

LuxDL · Sep 15, 2024 · df2bfd5 · df2bfd5
1 parent 4541d4f
commit df2bfd5
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/src/impl/matmul.jl b/src/impl/matmul.jl
@@ -244,7 +244,7 @@ function EnzymeRules.augmented_primal(cfg, ::EnzymeCore.Const{typeof(matmuladd!)
               !(C isa EnzymeCore.Const) ? copy(B.val) : nothing
 
     if !(C isa EnzymeCore.DuplicatedNoNeed || C isa EnzymeCore.BatchDuplicatedNoNeed)
-        matmuladd!(C.val, A.val, B.val, bias.val)
+        matmuladd!(C.val, opmode.val, A.val, B.val, bias.val)
     end
 
     return EnzymeRules.AugmentedReturn(nothing, nothing, (A_cache, B_cache))

diff --git a/test/common_ops/dense_tests.jl b/test/common_ops/dense_tests.jl
@@ -227,6 +227,9 @@ end
 
             (act === identity && hasbias) || continue
 
+            dweight .= 0
+            dx .= 0
+            db .= 0
             Enzyme.autodiff(Reverse, matmuladd!, Duplicated(y, copy(dy)),
                 Duplicated(weight, dweight), Duplicated(x, dx), b_enz)