Overload AbstractPPL.condition and AbstractPPL.decondition (#337)

Fixes #336.

Tests will fail until the Zygote bug is fixed... Maybe we should just mark them as broken so we can merge and release some PRs?

Co-authored-by: David Widmann <[email protected]>

Project.toml

@@ -1,6 +1,6 @@
 name = "DynamicPPL"
 uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8"
-version = "0.16.0"
+version = "0.16.1"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"

src/contexts.jl

@@ -366,11 +366,15 @@ otherwise return `context` which is [`DefaultContext`](@ref) by default.
 
 See also: [`decondition`](@ref)
 """
-condition(; values...) = condition(DefaultContext(), NamedTuple(values))
-condition(values::NamedTuple) = condition(DefaultContext(), values)
-condition(context::AbstractContext, values::NamedTuple{()}) = context
-condition(context::AbstractContext, values::NamedTuple) = ConditionContext(values, context)
-condition(context::AbstractContext; values...) = condition(context, NamedTuple(values))
+AbstractPPL.condition(; values...) = condition(DefaultContext(), NamedTuple(values))
+AbstractPPL.condition(values::NamedTuple) = condition(DefaultContext(), values)
+AbstractPPL.condition(context::AbstractContext, values::NamedTuple{()}) = context
+function AbstractPPL.condition(context::AbstractContext, values::NamedTuple)
+    return ConditionContext(values, context)
+end
+function AbstractPPL.condition(context::AbstractContext; values...)
+    return condition(context, NamedTuple(values))
+end
 
 """
     decondition(context::AbstractContext, syms...)
@@ -381,20 +385,22 @@ Note that this recursively traverses contexts, deconditioning all along the way.
 
 See also: [`condition`](@ref)
 """
-decondition(::IsLeaf, context, args...) = context
-function decondition(::IsParent, context, args...)
+AbstractPPL.decondition(::IsLeaf, context, args...) = context
+function AbstractPPL.decondition(::IsParent, context, args...)
     return setchildcontext(context, decondition(childcontext(context), args...))
 end
-decondition(context, args...) = decondition(NodeTrait(context), context, args...)
-function decondition(context::ConditionContext)
+function AbstractPPL.decondition(context, args...)
+    return decondition(NodeTrait(context), context, args...)
+end
+function AbstractPPL.decondition(context::ConditionContext)
     return decondition(childcontext(context))
 end
-function decondition(context::ConditionContext, sym)
+function AbstractPPL.decondition(context::ConditionContext, sym)
     return condition(
         decondition(childcontext(context), sym), BangBang.delete!!(context.values, sym)
     )
 end
-function decondition(context::ConditionContext, sym, syms...)
+function AbstractPPL.decondition(context::ConditionContext, sym, syms...)
     return decondition(
         condition(
             decondition(childcontext(context), syms...),

src/model.jl

@@ -254,8 +254,8 @@ From this we can tell what the correct way to condition `m` within `demo_inner`
 is in the two different models.
 
 """
-condition(model::Model; values...) = condition(model, NamedTuple(values))
-function condition(model::Model, values)
+AbstractPPL.condition(model::Model; values...) = condition(model, NamedTuple(values))
+function AbstractPPL.condition(model::Model, values)
     return contextualize(model, condition(model.context, values))
 end
 
@@ -307,7 +307,7 @@ julia> model(rng)
 (m = 0.683947930996541, x = 10.0)
 ```
 """
-function decondition(model::Model, syms...)
+function AbstractPPL.decondition(model::Model, syms...)
     return contextualize(model, decondition(model.context, syms...))
 end
 

test/compiler.jl

@@ -534,7 +534,7 @@ end
 
         # Ensure we can specialize on arguments.
         @model demo(x) = x ~ Normal()
-        length(methods(demo))
+        @test length(methods(demo)) == 4
         @test f(demo(1.0))
         f(::Model{typeof(demo),(:x,)}) = false
         @test !f(demo(1.0))

test/test_util.jl

@@ -43,7 +43,12 @@ function test_model_ad(model, logp_manual)
 
     y, back = Zygote.pullback(logp_model, x)
     @test y ≈ lp
-    @test back(1)[1] ≈ grad
+    # will be fixed by https://github.com/FluxML/Zygote.jl/pull/1104
+    if Threads.nthreads() > 1
+        @test_broken back(1)[1] ≈ grad
+    else
+        @test back(1)[1] ≈ grad
+    end
 end
 
 """