Error in `discard` returned by `update` for DynamicDSLTrace with hierarchical addresses

[fsaad]

Consider the following simple programs

using Gen

@gen function model_ok()
    k ~ poisson(5)
    for i=1:k
        {i} ~ uniform(0,1)
    end
end

@gen function model_bad()
    k ~ poisson(5)
    for i=1:k
        {:value => i} ~ uniform(0,1)
    end
end

Do we expect that Gen.update operates differently in these two cases. For model_ok, if we call update in such a way that reduces k then the discard address {i} are correctly placed in the discard, but in model_bad, the address {:value => i} are not in the discard. See following example:

Discard for model_ok

julia> tr, = Gen.generate(model_ok, (), choicemap(:k=>3));
julia> (new_trace, weight, _, discard) = Gen.update(tr, Gen.choicemap(:k=>1));
julia> display(discard)
│
├── :k : 3
│
├── 2 : 0.552591934644268
│
└── 3 : 0.6198635352707209

Discard for model_bad

julia> tr, = Gen.generate(model_bad, (), choicemap(:k=>3));
julia> (new_trace, weight, _, discard) = Gen.update(tr, Gen.choicemap(:k=>1));
julia> display(discard)
│
└── :k : 3

Possible Reason

The following lines seem to be the culprit: the recursive call to add_unvisited_to_discard! use an anonymous choicemap when subdiscard is empty, so the call to set_submap! fills out a subdiscard that we never access again.

Gen.jl/src/dynamic/update.jl

Lines 177 to 181 in 7955b07

	subdiscard = get_submap(discard, key)
	add_unvisited_to_discard!(
	isempty(subdiscard) ? choicemap() : subdiscard,
	subvisited, submap)
	set_submap!(discard, key, subdiscard)

Changing these lines to

                subdiscard = get_submap(discard, key)
                subdiscard_recursive = isempty(subdiscard) ? choicemap() : subdiscard
                add_unvisited_to_discard!(
                    subdiscard_recursive,
                    subvisited, submap)
                set_submap!(discard, key, subdiscard_recursive)

seems to fix the issue.

[fsaad]

Related #506

[ztangent]

Thanks for catching this!

So I'm trying to understand why the line of code you identified was introduced in the first place, and it looks like @alex-lew added it in this old commit: c8ce0d7

Apparently it was introduced to address cases where things were being discarded, but shouldn't be. So maybe we can check if the fix you suggested @fsaad also passes the original test case that @alex-lew added?

This is the test case:

Gen.jl/test/dsl/dynamic_dsl.jl

Lines 153 to 183 in 0570918

	# Addresses under the :data namespace will be visited,
	# but nothing there will be discarded.
	@gen function loopy()
	a = @trace(normal(0, 1), :a)
	for i=1:5
	@trace(normal(a, 1), :data => i)
	end
	end
	# Get an initial trace
	constraints = choicemap()
	constraints[:a] = 0
	for i=1:5
	constraints[:data => i] = 0
	end
	(trace,) = generate(loopy, (), constraints)
	# Update a
	constraints = choicemap()
	constraints[:a] = 1
	(new_trace, weight, retdiff, discard) = update(trace,
	(), (), constraints)
	# Test discard, score, weight, retdiff
	@test get_value(discard, :a) == 0
	prev_score = logpdf(normal, 0, 0, 1) * 6
	expected_new_score = logpdf(normal, 1, 0, 1) + 5 * logpdf(normal, 0, 1, 1)
	expected_weight = expected_new_score - prev_score
	@test isapprox(expected_new_score, get_score(new_trace))
	@test isapprox(expected_weight, weight)
	@test retdiff === UnknownChange()

[ztangent]

This issue also seems highly related, but @alex-lew identified a different line of code in add_unvisited_to_discard! as the likely culprit, so maybe it's a different issue?

#237

[mlb2251]

@yifr and I just ran into this same issue. We tried out @fsaad 's solution and it works.

We initially also thought it was the code @alex-lew points to, but actually that branch doesn't get run in @fsaad 's (and our own) examples because key in visited is only true when the hierarchical address comes from a GenerativeFunction call and not when you just define one like {:value => i} ~ ....

@fsaad 's reasoning for it being wrong is right – the original code has

add_unvisited_to_discard!( 
     isempty(subdiscard) ? choicemap() : subdiscard, 
     subvisited, submap)

But since add_unvisited_to_discard! actually works by mutating its first argument, it doesn't make sense to pass in a fresh choicemap without first keeping a reference to it as they do in their fix

I'll make a quick PR with the fix!