Update ADTypeCheckContext

Author: penelopeysm

test/ad.jl

@@ -137,14 +137,25 @@ 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)
+    # If we are using InitFromPrior or InitFromUniform to generate new values,
+    # then the parameter type will be Any, so we should skip the check.
+    lc = DynamicPPL.leafcontext(context)
+    if lc isa DynamicPPL.InitContext{
+        <:Any,<:Union{DynamicPPL.InitFromPrior,DynamicPPL.InitFromUniform}
+    }
+        return nothing
+    end
+    # 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 !(any(param_eltype .<: valids))
+        throw(IncompatibleADTypeError(param_eltype, adtype(context)))
     end
-    return nothing
 end
 
 # A bunch of tilde_assume/tilde_observe methods that just call the same method on the child
@@ -199,10 +210,10 @@ end
             @testset "Expected: $expected_adtype, Actual: $actual_adtype" begin
                 if actual_adtype == expected_adtype
                     # Check that this does not throw an error.
-                    sample(contextualised_tm, sampler, 2)
+                    sample(contextualised_tm, sampler, 2; check_model=false)
                 else
                     @test_throws AbstractWrongADBackendError sample(
-                        contextualised_tm, sampler, 2
+                        contextualised_tm, sampler, 2; check_model=false
                     )
                 end
             end