Fix primal again

Project.toml

@@ -1,7 +1,7 @@
 name = "Checkpointing"
 uuid = "eb46d486-4f9c-4c3d-b445-a617f2a2f1ca"
 authors = ["Michel Schanen <[email protected]>", "Sri Hari Krishna Narayanan <[email protected]>"]
-version = "0.9.2"
+version = "0.9.3"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/Rules/EnzymeRules.jl

@@ -11,11 +11,12 @@ function augmented_primal(
     model,
     range,
 )
-    primal = func.val(body.val, alg.val, deepcopy(model.val), range.val)
+    tape_model = deepcopy(model.val)
+    func.val(body.val, alg.val, model.val, range.val)
     if needs_primal(config)
-        return AugmentedReturn(primal, nothing, (model.val,))
+        return AugmentedReturn(nothing, nothing, (tape_model,))
     else
-        return AugmentedReturn(nothing, nothing, (model.val,))
+        return AugmentedReturn(nothing, nothing, (tape_model,))
     end
 end
 
@@ -30,14 +31,13 @@ function reverse(
     range,
 )
     (model_input,) = tape
-    model_final = Checkpointing.rev_checkpoint_struct_for(
+    Checkpointing.rev_checkpoint_struct_for(
         body.val,
         alg.val,
         model_input,
         model.dval,
         range.val,
     )
-    copyto!(model.val, model_final)
     return (nothing, nothing, nothing, nothing)
 end
 
@@ -50,11 +50,12 @@ function augmented_primal(
     model,
     condition,
 )
-    primal = func.val(body.val, alg.val, deepcopy(model.val), condition.val)
+    tape_model = deepcopy(model.val)
+    func.val(body.val, alg.val, model.val, condition.val)
     if needs_primal(config)
-        return AugmentedReturn(primal, nothing, (model.val,))
+        return AugmentedReturn(nothing, nothing, (tape_model,))
     else
-        return AugmentedReturn(nothing, nothing, (model.val,))
+        return AugmentedReturn(nothing, nothing, (tape_model,))
     end
 end
 
@@ -69,14 +70,13 @@ function reverse(
     condition,
 )
     (model_input,) = tape
-    model_final = Checkpointing.rev_checkpoint_struct_while(
+    Checkpointing.rev_checkpoint_struct_while(
         body.val,
         alg.val,
         model_input,
         model.dval,
         condition.val,
     )
-    copyto!(model.val, model_final)
     return (nothing, nothing, nothing, nothing)
 end
 

test/multilevel.jl

@@ -43,5 +43,5 @@ g = autodiff(
 )
 
 # TODO: Primal is wrong only when multilevel checkpointing is used
-@test_broken g[2] == primal
+@test g[2] == primal
 @test all(dx.x .== [1024.0, 1024.0, 1024.0])