use special construct to hijack assume for :=

TuringLang · Apr 25, 2024 · b1c3206 · b1c3206
1 parent 503d4de
commit b1c3206
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 3 deletions.
diff --git a/src/compiler.jl b/src/compiler.jl
@@ -387,7 +387,6 @@ function generate_mainbody!(mod, found, expr::Expr, warn)
 end
 
 function generate_assign(left, right)
-    right_expr = :($(Distributions.Dirac)($right))
     tilde_expr = generate_tilde(left, right_expr)
     return quote
         if $(is_extracting_values)(__context__)

diff --git a/src/values_as_in_model.jl b/src/values_as_in_model.jl
@@ -1,3 +1,11 @@
+struct TrackedValue{T}
+    value::T
+end
+
+is_tracked_value(::TrackedValue) = true
+is_tracked_value(::Any) = false
+
+check_tilde_rhs(x::TrackedValue) = x
 
 """
     ValuesAsInModelContext
@@ -55,7 +63,12 @@ end
 
 # `tilde_asssume`
 function tilde_assume(context::ValuesAsInModelContext, right, vn, vi)
-    value, logp, vi = tilde_assume(childcontext(context), right, vn, vi)
+    if is_tracked_value(right)
+        value = right.value
+        logp = zero(getlogp(vi))
+    else
+        value, logp, vi = tilde_assume(childcontext(context), right, vn, vi)
+    end
     # Save the value.
     push!(context, vn, value)
     # Save the value.
@@ -65,7 +78,12 @@ end
 function tilde_assume(
     rng::Random.AbstractRNG, context::ValuesAsInModelContext, sampler, right, vn, vi
 )
-    value, logp, vi = tilde_assume(rng, childcontext(context), sampler, right, vn, vi)
+    if is_tracked_value(right)
+        value = right.value
+        logp = zero(getlogp(vi))
+    else
+        value, logp, vi = tilde_assume(rng, childcontext(context), sampler, right, vn, vi)
+    end
     # Save the value.
     push!(context, vn, value)
     # Pass on.