Merge branch 'master' into torfjelde/transformations

TuringLang · Feb 3, 2024 · b0dd2f8 · b0dd2f8
2 parents d7aba55 + c33eeae
commit b0dd2f8
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/src/context_implementations.jl b/src/context_implementations.jl
@@ -16,10 +16,22 @@ require_particles(spl::Sampler) = false
 
 # Allows samplers, etc. to hook into the final logp accumulation in the tilde-pipeline.
 function acclogp_assume!!(context::AbstractContext, vi::AbstractVarInfo, logp)
+    return acclogp_assume!!(NodeTrait(acclogp_assume!!, context), context, vi, logp)
+end
+function acclogp_assume!!(::IsParent, context::AbstractContext, vi::AbstractVarInfo, logp)
+    return acclogp_assume!!(childcontext(context), vi, logp)
+end
+function acclogp_assume!!(::IsLeaf, context::AbstractContext, vi::AbstractVarInfo, logp)
     return acclogp!!(context, vi, logp)
 end
 
 function acclogp_observe!!(context::AbstractContext, vi::AbstractVarInfo, logp)
+    return acclogp_observe!!(NodeTrait(acclogp_observe!!, context), context, vi, logp)
+end
+function acclogp_observe!!(::IsParent, context::AbstractContext, vi::AbstractVarInfo, logp)
+    return acclogp_observe!!(childcontext(context), vi, logp)
+end
+function acclogp_observe!!(::IsLeaf, context::AbstractContext, vi::AbstractVarInfo, logp)
     return acclogp!!(context, vi, logp)
 end