Fix MLE/MAP with Zygote and ReverseDiff (#1408)

* Fix AD for modes

* Increment version number

* Remove display calls

* Move tests

Project.toml

@@ -1,6 +1,6 @@
 name = "Turing"
 uuid = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
-version = "0.14.2"
+version = "0.14.3"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"

src/core/ad.jl

@@ -65,7 +65,8 @@ getADbackend(spl::SampleFromPrior) = ADBackend()()
         θ::AbstractVector{<:Real},
         vi::VarInfo,
         model::Model,
-        sampler::AbstractSampler=SampleFromPrior(),
+        sampler::AbstractSampler,
+        ctx::DynamicPPL.AbstractContext = DynamicPPL.DefaultContext()
     )
 
 Computes the value of the log joint of `θ` and its gradient for the model
@@ -89,6 +90,7 @@ gradient_logp(
     vi::VarInfo,
     model::Model,
     sampler::AbstractSampler = SampleFromPrior(),
+    ctx::DynamicPPL.AbstractContext = DynamicPPL.DefaultContext()
 )
 
 Compute the value of the log joint of `θ` and its gradient for the model
@@ -160,7 +162,7 @@ function gradient_logp(
     # Specify objective function.
     function f(θ)
         new_vi = VarInfo(vi, sampler, θ)
-        model(new_vi, sampler)
+        model(new_vi, sampler, context)
         return getlogp(new_vi)
     end
 

src/core/compat/reversediff.jl

@@ -27,7 +27,7 @@ function gradient_logp(
     # Specify objective function.
     function f(θ)
         new_vi = VarInfo(vi, sampler, θ)
-        model(new_vi, sampler)
+        model(new_vi, sampler, context)
         return getlogp(new_vi)
     end
     tp, result = taperesult(f, θ)
@@ -65,7 +65,7 @@ end
         # Specify objective function.
         function f(θ)
             new_vi = VarInfo(vi, sampler, θ)
-            model(new_vi, sampler)
+            model(new_vi, sampler, context)
             return getlogp(new_vi)
         end
         ctp, result = memoized_taperesult(f, θ)

test/modes/ModeEstimation.jl

@@ -6,6 +6,7 @@ using NamedArrays
 using ReverseDiff
 using Random
 using LinearAlgebra
+using Zygote
 
 dir = splitdir(splitdir(pathof(Turing))[1])[1]
 include(dir*"/test/test_utils/AllUtils.jl")

test/runtests.jl

@@ -40,6 +40,10 @@ include("test_utils/AllUtils.jl")
         @testset "variational algorithms : $adbackend" begin
             include("variational/advi.jl")
         end
+
+        @testset "modes" begin
+            include("modes/ModeEstimation.jl")
+        end
     end
     @testset "variational optimisers" begin
         include("variational/optimisers.jl")
@@ -55,8 +59,4 @@ include("test_utils/AllUtils.jl")
         # include("utilities/stan-interface.jl")
         include("inference/utilities.jl")
     end
-
-    @testset "modes" begin
-        include("modes/ModeEstimation.jl")
-    end
 end