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