Check ADType use in optimisation

Author: mhauru

test/optimisation/Optimisation.jl

@@ -1,6 +1,7 @@
 module OptimisationTests
 
 using ..Models: gdemo, gdemo_default
+using ..ADUtils: ADTypeCheckContext
 using Distributions
 using Distributions.FillArrays: Zeros
 using DynamicPPL: DynamicPPL
@@ -140,7 +141,6 @@ using Turing
                 gdemo_default, OptimizationOptimJL.LBFGS(); initial_params=true_value
             )
             m3 = maximum_likelihood(gdemo_default, OptimizationOptimJL.Newton())
-            # TODO(mhauru) How can we check that the adtype is actually AutoReverseDiff?
             m4 = maximum_likelihood(
                 gdemo_default, OptimizationOptimJL.BFGS(); adtype=AutoReverseDiff()
             )
@@ -616,6 +616,17 @@ using Turing
         @assert vcat(get_a[:a], get_b[:b]) == result.values.array
         @assert get(result, :c) == (; :c => Array{Float64}[])
     end
+
+    @testset "ADType" begin
+        Random.seed!(222)
+        for adbackend in (AutoReverseDiff(), AutoForwardDiff(), AutoTracker())
+            m = DynamicPPL.contextualize(
+                gdemo_default, ADTypeCheckContext(adbackend, gdemo_default.context)
+            )
+            maximum_likelihood(m; adtype=adbackend)
+            maximum_a_posteriori(m; adtype=adbackend)
+        end
+    end
 end
 
 end

test/runtests.jl

@@ -7,6 +7,7 @@ import Turing
 
 include(pkgdir(Turing) * "/test/test_utils/models.jl")
 include(pkgdir(Turing) * "/test/test_utils/numerical_tests.jl")
+include(pkgdir(Turing) * "/test/test_utils/ad_utils.jl")
 
 Turing.setprogress!(false)
 

test/test_utils/ad_utils.jl

@@ -96,7 +96,7 @@ end
 """
     valid_eltypes(context::ADTypeCheckContext)
 
-Return the element types that are valid for the ADType of `context`.
+Return the element types that are valid for the ADType of `context` as a tuple.
 """
 function valid_eltypes(context::ADTypeCheckContext)
     context_at = adtype(context)
@@ -116,7 +116,7 @@ Check that the element types in `vi` are compatible with the ADType of `context`
 
 Throw an `IncompatibleADTypeError` if an incompatible element type is encountered.
 """
-function check_adtype(context::ADTypeCheckContext, vi::DynamicPPL.VarInfo)
+function check_adtype(context::ADTypeCheckContext, vi::DynamicPPL.AbstractVarInfo)
     valids = valid_eltypes(context)
     for val in vi[:]
         valtype = typeof(val)
@@ -187,6 +187,14 @@ function DynamicPPL.dot_tilde_observe(context::ADTypeCheckContext, right, left,
     return logp, vi
 end
 
+function DynamicPPL.dot_tilde_observe(context::ADTypeCheckContext, sampler, right, left, vi)
+    logp, vi = DynamicPPL.dot_tilde_observe(
+        DynamicPPL.childcontext(context), sampler, right, left, vi
+    )
+    check_adtype(context, vi)
+    return logp, vi
+end
+
 # Check that the ADTypeCheckContext works as expected.
 Test.@testset "ADTypeCheckContext" begin
     Turing.@model test_model() = x ~ Turing.Normal(0, 1)