Open
Description
Yesterday @alex-lew and I found a bug in the implementation of add_unvisited_to_discard! on the 20220821-prox-base branch. The bug can be reproduced by instantiating the InverseGraphics project with Gen coming from the 20220821-prox-base branch and replacing the likelihood call with a call to a Prox version of the likelihood (e.g. the one found here). The bug occurs when running line 139 of the demo.jl notebook.
The bug seems to come from the incorrect handling of visited submaps of type ValueChoiceMap. Alex and I came up with the following fix. I have not reasoned carefully about the fix so I'm not sure it is a complete fix. I also don't think it is the most elegant fix, but I'm including it for reference.
function add_unvisited_to_discard!(discard::DynamicChoiceMap,
visited::DynamicSelection,
prev_choices::ChoiceMap)
for (key, value) in get_values_shallow(prev_choices)
if !(key in visited)
@assert !has_value(discard, key)
@assert isempty(get_submap(discard, key))
set_value!(discard, key, value)
end
end
for (key, submap) in get_submaps_shallow(prev_choices)
if submap isa ValueChoiceMap
continue
end
@assert !has_value(discard, key)
if key in visited
# the recursive call to update already handled the discard
# for this entire submap
continue
else
subvisited = visited[key]
if isempty(subvisited)
# none of this submap was visited, so we discard the whole thing
@assert isempty(get_submap(discard, key))
set_submap!(discard, key, submap)
else
subdiscard = get_submap(discard, key)
add_unvisited_to_discard!(
isempty(subdiscard) ? choicemap() : subdiscard,
subvisited, submap)
set_submap!(discard, key, subdiscard)
end
end
end
end