Discrete single-site map optimize and gibbs sampling

It's easy to implement single-site versions of these, where the user provides a (finite) set of values of the variable to consider. These would go into the inference library:

Map optimization:

function Gen.map_optimize(trace, addr, values)
    args = get_args(trace)
    argdiffs = map((_) -> NoChange(), args)
    weights = Vector{Float64}(undef, length(values))
    traces = Vector{Any}(undef, length(values))
    for (i, val) in enumerate(values)
        (traces[i], weights[i], _, _) = update(trace, args, argdiffs, choicemap((addr, val)))
    end
    idx = argmax(weights)
    return traces[idx]
end

Computing probabilities for Gibbs sampling:

function Gen.discrete_conditional_dist(trace, addr, values)
    args = get_args(trace)
    argdiffs = map((_) -> NoChange(), args)
    weights = Vector{Float64}(undef, length(values))
    for (i, val) in enumerate(values)
        (_, weights[i]) = update(trace, args, argdiffs, choicemap((addr, val)))
    end
    exp.(weights .- logsumexp(weights))
end

Then, a built-in Gibbs sampling procedure would be:

function Gen.discrete_conditional_update(trace, addr, values)
    probs = discrete_conditional_dist(trace, addr, values)
    idx = uniform_discrete(probs)
    val = values[idx]
    args = get_args(trace)
    argdiffs = map((_) -> NoChange(), args)
    new_trace, = update(trace, args, argdiffs, choicemap((addr, val)))
    return new_trace
end

Note that both the MAP and Gibbs sampling procedures assume that the set of addresses in the trace (i.e. control flow) is not affected by the choice being changed. This is an important requirement, as doing "Gibbs" in the presence of stochastic structure would need to be carefully defined, and could end up requiring auxiliary-variable based approximation (note: I think that the old fix-update could be useful for that).

Alternatively, discrete conditional probabilities (from discrete_conditional_dist) could be used directly in a custom proposal (which is useful for block proposals when some of the variables are proposed using Gibbs updates, conditioned on the value proposed for earlier variables).

For best performance, generating a static choice map just once per call, or statically, should be considered.