Update ADTypeCheckContext

Author: penelopeysm

test/ad.jl

@@ -137,14 +137,17 @@ 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.AbstractVarInfo)
+    # Note that `get_param_eltype` will return `Any` with e.g. InitFromPrior or
+    # InitFromUniform, so this will fail. But on the bright side, you would never _really_
+    # use AD with those strategies, so that's fine. The cases where you do want to
+    # use this are DefaultContext (i.e., old, slow, LogDensityFunction) and
+    # InitFromParams{<:VectorWithRanges} (i.e., new, fast, LogDensityFunction), and
+    # both of those give you sensible results for `get_param_eltype`.
+    param_eltype = DynamicPPL.get_param_eltype(vi, context)
     valids = valid_eltypes(context)
-    for val in vi[:]
-        valtype = typeof(val)
-        if !any(valtype .<: valids)
-            throw(IncompatibleADTypeError(valtype, adtype(context)))
-        end
+    if !(param_eltype .<: valids)
+        throw(IncompatibleADTypeError(valtype, adtype(context)))
     end
-    return nothing
 end
 
 # A bunch of tilde_assume/tilde_observe methods that just call the same method on the child
@@ -176,7 +179,11 @@ end
 """
 All the ADTypes on which we want to run the tests.
 """
-ADTYPES = [AutoForwardDiff(), AutoReverseDiff(; compile=false)]
+adtypes = (
+    AutoForwardDiff(),
+    AutoReverseDiff(),
+    # Don't need to test Mooncake as it doesn't use tracer types
+)
 if INCLUDE_MOONCAKE
     push!(ADTYPES, AutoMooncake(; config=nothing))
 end