Fix a bug in nested decondition_context

TuringLang · Jan 9, 2025 · 7161599 · 7161599
1 parent 90254fa
commit 7161599
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 3 deletions.
diff --git a/src/contexts.jl b/src/contexts.jl
@@ -447,7 +447,9 @@ function decondition_context(context::ConditionContext, sym, syms...)
         # No more values left, can unwrap
         decondition_context(childcontext(context), syms...)
     else
-        ConditionContext(new_values, decondition_context(childcontext(context), syms...))
+        ConditionContext(
+            new_values, decondition_context(childcontext(context), sym, syms...)
+        )
     end
 end
 function decondition_context(context::NamedConditionContext, vn::VarName{sym}) where {sym}

diff --git a/test/contexts.jl b/test/contexts.jl
@@ -133,7 +133,7 @@ end
                         @test hasconditioned_nested(context, vn_child)
                         # Value should be the same as extracted above.
                         @test getconditioned_nested(context, vn_child) ===
-                            getoptic(vn_child)(val)
+                              getoptic(vn_child)(val)
                     end
                 end
             end
@@ -257,7 +257,25 @@ end
                 @test dctx.values == Dict(@varname(x) => 1)
                 # Decondition all variables manually
                 @test decondition_context(ctx, @varname(x), @varname(y), @varname(z)) isa
-                    DefaultContext
+                      DefaultContext
+            end
+
+            @testset "Nesting" begin
+                ctx = ConditionContext(
+                    (x=1, y=2),
+                    ConditionContext(Dict(@varname(a) => 3, @varname(b) => 4))
+                )
+                # Decondition an outer variable
+                dctx = decondition_context(ctx, :x)
+                @test dctx.values == (y=2,)
+                @test childcontext(dctx).values == Dict(@varname(a) => 3, @varname(b) => 4)
+                # Decondition an inner variable
+                dctx = decondition_context(ctx, @varname(a))
+                @test dctx.values == (x=1, y=2)
+                @test childcontext(dctx).values == Dict(@varname(b) => 4)
+                # Try deconditioning everything
+                dctx = decondition_context(ctx)
+                @test dctx isa DefaultContext
             end
         end
     end