Incorrect Dynamic DSL update Implementation on the Prox Base Branch

[mugamma]

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

[mugamma]

@georgematheos @femtomc Tagging you here as I think you might have a better idea for the fix.

[femtomc]

@mugamma Can you provide a bit more context about the issue:

• Is discard missing choices which should be discarded according to the visited?
• Did you notice this because of an incorrect weight computation (due to the above)?
• Can you provide more context about force_structure? Why does the bug occur here -- what is force_structure doing so that you noticed the bug?

E.g. what exactly is the bug -- is there a stacktrace? Or is it a math semantics bug (like the above)?

[femtomc]

I think I can guess the issue here:

• ValueChoiceMap doesn't have internal submaps, so get_submaps_shallow(vchm::ValueChoiceMap) returns an empty iterator.
• Okay, but it does have a value -- so what does get_values_shallow(vchm::ValueChoiceMap) return? For get_values_shallow, the assumption is that the implementation returns an iterator of (key, val) pairs.
• I don't know what get_values_shallow on ValueChoiceMap returns -- but I'm assuming the key is either a constant literal or nothing.

One way I fixed this in GenJAX is that I don't have a get_values_shallow - I only have a get_submaps_shallow and a get_leaf_value. (Someone can correct me if I'm wrong) but my reasoning was that get_values_shallow seems like a vestige of the separation between generative functions and distributions.

To me, "the right fix" is just to eliminate get_values_shallow and have two types of ChoiceMap:

• "Leaf" ChoiceMap (e.g. ValueChoiceMap)
• Non-leaf ChoiceMap (e.g. DynamicChoiceMap)

In lieu of correcting the type/interface hierarchy, if the implementation you posted is "working" now -- I'm concerned about how ValueChoiceMap participates in get_values_shallow and get_submaps_shallow -- if you continue when you get a ValueChoiceMap in a pair (key, chm) from get_submaps_shallow, aren't you potentially missing those choices now in discard?

Also, I don't understand these lines:

        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

You check if key is in visited, and then in the else branch, you get subvisited = visited[key] -- won't that always be empty?

[femtomc]

Right, so -- for ValueChoiceMap, here's what the implementation of get_values_shallow on the branch you linked implies:

# get_values_shallow and get_nonvalue_submaps_shallow are just filters on get_submaps_shallow
"""
    get_values_shallow(choices::ChoiceMap)

Returns an iterable collection of tuples `(address, value)`
for each value stored at a top-level address in `choices`.
(Works by applying a filter to `get_submaps_shallow`,
so this internally requires iterating over every submap.)
"""
function get_values_shallow(choices::ChoiceMap)
    (
        (addr, get_value(submap))
        for (addr, submap) in get_submaps_shallow(choices)
        if has_value(submap)
    )
end

This returns an empty iterator for ValueChoiceMap because of

Gen.jl/src/choice_map/choice_map.jl

Line 152 in af6fbfc

@inline get_submaps_shallow(choices::ValueChoiceMap) = ()

.

@mugamma It seems like, for that branch (which includes the distribution as generative function changes) -- the choice map interfaces (implemented for ValueChoiceMap) are the following:

@inline has_value(choices::ValueChoiceMap) = true
@inline get_value(choices::ValueChoiceMap) = choices.val
@inline get_submap(choices::ValueChoiceMap, addr) = EmptyChoiceMap()
@inline get_submaps_shallow(choices::ValueChoiceMap) = ()
@inline Base.:(==)(a::ValueChoiceMap, b::ValueChoiceMap) = a.val == b.val
@inline Base.:(==)(a::ValueChoiceMap, b::ChoiceMap) = false
@inline Base.:(==)(a::ChoiceMap, b::ValueChoiceMap) = false
@inline Base.isapprox(a::ValueChoiceMap, b::ValueChoiceMap) = isapprox(a.val, b.val)
@inline get_address_schema(::Type{<:ValueChoiceMap}) = AllAddressSchema()

And get_values_shallow just uses the new interface functions -- so I'm guessing that the bug occurs because something isn't quite working correctly (e.g. it didn't get the right changes) for the new interfaces.

perhaps the right fix is just to use get_submaps_shallow, and the has_value/get_value interfaces in the implementation of add_unvisited_to_discard! above.

[alex-lew]

@femtomc @mugamma I thought the issue didn't have to do with ValueChoiceMaps. The reason it shows up here is because @georgematheos -- after introducing ValueChoiceMaps -- simplified the code to handle values and submaps in a uniform way, but this ran into another bug, having to do with selections.

The issue arises when a dynamic DSL generative function does something like {:x => :y} ~ normal(0, 1). Note that this creates a nested ChoiceMap despite the fact that there is no call for the address :x.

As a result, when we call update on this generative function (where :x => :y is one of the updated addresses):

• visited contains the address :x => :y, but does not contain the address :x, because :x itself was not "visited" by the update call.
• on L100, we iterate over the submaps in the previous trace, one of which is at address :x
• on L105, we trigger the assertion, because visisted doesn't have :x but the discard is non-empty at :x (it has a value at :y in the submap).

So, in response to @femtomc's question:

You check if key is in visited, and then in the else branch, you get subvisited = visited[key] -- won't that always be empty?

No -- if I write visited = select(:x => :y), then !(:x in visited) but visited[:x] is select(:y).

[femtomc]

Thanks Alex for clarifying!

Dumb suggestion — but can you implement this function by taking the complement of the visited selection? And then applying this complement to filter the previous choice map and merging the result with the discard created from the update visitation?

[fsaad]

Related #512