(Ready for review): Switch combinator

[femtomc]

Beginning work on Switch combinator. Initial version of propose and generate. Beginning testing.

[femtomc]

My hope is that a few people can take a look at the trace, and the sketch up for propose and generate and provide feedback. As per the discussion with Marco on Slack:

Right. We could imagine various versions, from the least restrictive (i) which accepts a Bool argument that decides which branch to choose, to (ii) a version that calls a generative function to determine the branch as you wrote above, to (iii) a version that accepts a probability argument and samples the branch from a Bernoulli as in our trace types paper. I think the choice of which one to implement depends on what analyses they make easier. My hunch is that either the first one (which permits straightforward enumeration of trace structures i.e. data-flow paths) or the last one as in the trace types paper, would be the most useful. You could also imagine transformations that lower from (iii) to (i) -- which is straightforward, or analysis that tries to do pattern matching to go from (iii) to (i) -- which is harder. If I had to guess, I would recommend going with (iii), but making that decision confidently would seem to depend on having a few analyses in mind. There is a tradeoff between (a) expressiveness and (b) exposing high-level information relevant for analysis. (i) is more expressive than (iii) but provides less information to analyses

My proposal will follow the outline of the Switch combinator expressed in Trace types and denotational semantics:

Note that the following restrictions apply:

1. Branches must have the same trace type. This may be a conservative constraint - but it seemed natural.
2. Each branch must accept the same argument tuple type. This may also be conservative - but it does lead to stronger static properties. In the future, I can imagine lifting this restriction. My main concern here was argument difference types - it seemed simpler to reason about the same set of arguments.

[marcoct]

@femtomc This samples from a Bernoulli distribution with probability branch_p and returns a Bool.

[femtomc]

Apologies - that comment needs to be removed. I figured that bit out and corrected it elsewhere, but the comments indicate otherwise.

[marcoct]

@femtomc This is exciting! A few notes:

I think WithProbability (using the name from the POPL paper) would probably be a good name for this combinator, so we can reserve Switch for the more generic combinator that takes a vector of generative functions and returns a generative function that accepts an integer argument that switches between them. (Cond is also a reasoanable name for that, I think).

@femtomc I thought more about the different combinators in this space:

I think the combinator that takes a vector of generative functions and returns a generative function that accepts an integer argument that switches between them (I'll call this Switch in this comment), is higher priority for the project as a whole, since it would bring the expressiveness of the SML + combinators language up to a reasonable stable point -- we have applications right now where we can't write a model in SML + combinators because we need to switch between several different branches, and having that feature would allow us to write those models in SML + combinators and would unlock specific features that are specific to SML + combinators for many applications, including the incremental computation capabilities and the variable elimination compiler I've been working on.

Of course, there are many possibilities:

1. if-else branching with bernoulli random choice in the trace
2. if-else branching with boolean condition argument
3. multiple-branch switching with categorical random choice in the trace
4. multiple-branch switching with int argument to decide what branch to run

From the POPL paper, the main benefit of WithProbability is that, if the probability is between 0 and 1 (exclusive of both), then we can guarantee that the density on traces is nonzero, which is useful for verifying certain relationships between proposals and models hold. But it is also possible to define a notion of a 'trace type' that doesn't guarantee nonzero density, but provides an over-approximation to the support of the density on traces that would still be practically useful. For option 1 and option 4 above, it would still make sense to define a trace type that is the union of the various trace types for each branch, and it would still be possible to staically find many bugs in practice using such a trace type analysis. For example, we could check statically that the trace types for a proposal and a model generative function have at least some traces in both (as opposed to proving that they are identical). That would catch some common errors, like forgetting to sample a choice in the proposal. But I think some more thinking on that would be needed.

In general, I think the number of combinators can grow, and we it will take some experience using them to understand which combinators are most useful, and whether some tweaks would be useful.

Some specific notes about the combinator you've started to implement:

Branches must have the same trace type. This may be a conservative constraint - but it seemed natural.

By trace type you mean the set of possible choice maps it samples, right? I think this is overly restrictive. We want to be able to handle cases where the two branches make random choices at different addresses.

Each branch must accept the same argument tuple type. This may also be conservative - but it does lead to stronger static properties. In the future, I can imagine lifting this restriction. My main concern here was argument difference types - it seemed simpler to reason about the same set of arguments.

I think that's reasonable. Also, I think it makes sense also to require that the return types of the two generative functions are the same. Together those two restrictions will make things like AD more straightforward.

This combinator manifests some interesting design choices relating to the probabilistic semantics of generative functions. There are some design choices to be made for the addressing scheme. For example, the Bernoulli random choice needs to have an address of its own. Also, do addresses sampled in the two branches live in the same namespace or different namespaces? If they live in the same namespace, then the semantics of update and regenerate mean that these operations will need to automatically copy data from one branch's trace to another. The version where the branches have different namespaces for the two branches will be easier to implement, and I think that probably makes sense for a first version.

[femtomc]

@marcoct

By trace type you mean the set of possible choice maps it samples, right? I think this is overly restrictive. We want to be able to handle cases where the two branches make random choices at different addresses.

No, I actually mean the subtype of Trace which the user uses.

Sorry, trace type is now overloaded. I should use inheritor of Trace - that's what I mean here.

Specifically, all I'm saying is that you have to use the same generative function type in each branch. Do you think that is too conservative?

Edit: generative function type.

[alex-lew]

Specifically, all I'm saying is that you have to use the same generative function in each branch. Do you think that is too conservative?

Interesting -- why do we need a combinator to branch between applications of the same generative function? Can't we do

args = condition ? args1 : args2
x ~ f(args...)

in that case?

[femtomc]

@alex-lew Oops - generative function type. corrected.

[marcoct]

@femtomc Yes, I think it would be better to allow the subtraces to have any Julia type that is a subtype of Gen.Trace and for the generative functions to be of any Julia type that is a subtype of Gen.GenerativeFunction (but with the same return-value type). I think requiring that they have the same return-value type is useful.

[alex-lew]

I think requiring that they have the same return-value type is useful.

Presumably it's all right if this is technically a Union type :-) Or at least, we may need to handle the case where both GFs return union types, and it doesn't seem to be more challenging to handle the case where they return different things? But I might not be thinking through the subtleties with respect to AD.

[femtomc]

@marcoct I've started thinking about the difference between (1 and 2) and (3 and 4) a bit.

In practice, is there any difference when extracting the internal random choice made by the combinators described by (1, 3) into a choice preceding (2, 4) ?

Neither set of implementations appears to restrict the modeller. I also couldn't think of anything from an argdiff perspective. It seems mostly a matter of convenience (because we can isolate the "which branch" choice inside the namespace of the combinator, whereas the user might have to populate their toplevel model namespace with multiple branch choices).

Right now, I'm planning on implementing 1 and 4.

Also just a comment - I don't think there's a large increase in complexity for deriving the trace type for (2, 3) ? Because e.g. I would just type these as:

:outer address -> :cond -> Bernoulli
                  :branch -> Sum{TraceType G1, TraceType G2}

etc. But I might be missing something.

[femtomc]

@marcoct In terms of namespaces, the semantics as you mentioned above are really interesting. If we take the sharing route, what happens if the address spaces for either branch have non-empty set difference? When I switch - do I update with constraints on addresses in the intersection, even if I can't visit all the constraints in the switch?

Conversely, do I partially constrain and then sample from the prior?

I expect that both of these define a valid update if we implement the correct reverse move (e.g. re-score the constraints and subtract the previous branch weight) but this is different than the branch switching behavior in the dynamic DSL (iirc). But it's neat!

[femtomc]

@marcoct This seems correct for the "namespace merging" semantics. In update, if I switch branches, I generate with a merge of the previous traces choices and the choice map. I'm worried this will throw an error (however) if not all constraints are visited (e.g. if the previous trace has non-empty set difference with the new namespace).

I'm also dispatching on process! using the diff types. I'm not sure if this current implementation is correct - mostly because I'm not supported Diffed dispatch yet (and I suspect I have to).

[marcoct]

I'm also dispatching on process! using the diff types. I'm not sure if this current implementation is correct - mostly because I'm not supported Diffed dispatch yet (and I suspect I have to).

@femtomc The Diffed types are only used for boxing Julia values for diff propagation of Julia code via operator overloading. GFI implementers only need to worry about Diff values. So what you have is right.

[marcoct]

@marcoct In terms of namespaces, the semantics as you mentioned above are really interesting. If we take the sharing route, what happens if the address spaces for either branch have non-empty set difference? When I switch - do I update with constraints on addresses in the intersection, even if I can't visit all the constraints in the switch?

Given an addressing scheme, the semantics for update are well-defined -- the addresses that are shared get copied over, unless they are themselves constrained. New addresses that are introduced but not constrained get sampled from an internal proposal distribution. The discard contains all addresses that were removed, and any shared addresses that were constrained. The weight is the ratio of densities of the two traces, with an extra factor in the denominator for the density on the newly introduced choices that were not constrained. This behavior is described rather tersely in the docs here. An easier-to-read version of the semantics of update is given in page 71 of my thesis. (The definition in my thesis defines the behavior when there are newly introduced but unconstrained choices as an error, and the definition in the Gen documentation extends that definition to allow those choices get sampled from the internal from the internal proposal.).

I recognize that the version in the Gen docs at the moment is rather terse, so here is what I think the behavior should be:

In the version where the two branches use different address namespaces, when 'update' receives a change in the branch, it would call 'generate' -- with no constraints -- on the new branch to fill in the choices. In the version where the two branches use the same address namespace, you would extract the choices from the previous branch's trace and use them as constraints that you pass to a call to 'generate' on the new branch's generative function. I think in both cases the log weight returned by 'update' would be lw1 - lw2 where lw1 is the log-weight returned by 'generate' and where lw2 is the result of calling get_score on the previous branch's trace. In the first case, the discard is always empty, and in the second case the discard is merge of (i) the set of addresses that were in the first branch and not the second, and (ii) the set of addresses that are in both branches and also in the constraints that were passed to 'update'. The retdiff value would be UnknownChange.

I think it would be good to implement the version where they use different namespaces first -- the more complicated version can always be a separate PR. Later on, I could also imagine a flag that you pass to the combinator that specifies which version gets generated -- whether they use shared namespace or not -- and also for the different namespace version, what the addresses to use for each branch would be).

[femtomc]

@marcoct I think I have the shared address version implemented. I tested to make sure the (score - weight) after update is the same as the score before (to verify that the implementation was consistent). I suspect when I formalize my test file - I'll include tests which do this (just like tests for update and regenerate do now) as well.

The last thing I need to check off is the discard behavior.

[marcoct]

@marcoct This seems correct for the "namespace merging" semantics. In update, if I switch branches, I generate with a merge of the previous traces choices and the choice map. I'm worried this will throw an error (however) if not all constraints are visited (e.g. if the previous trace has non-empty set difference with the new namespace).

@femtomc Nice! Just read your comment above, and now I better understand your first question. Yes, I think we need to relax the requirement on generate so that it does not fail if there are unvisited choices. More specifically, there should be a flag that can be used to turn on or off this check. In my thesis I documented 'generate' as having this behavior, to support things like what we're trying to do here, and I think the Gen implementation should match this. (I think it already does, but it's not well documented.) I just made an issue for this: #335

[femtomc]

@marcoct I had to construct some custom methods to correctly perform the merging and discard computation.

One question I had:

function update_recurse_merge(prev_choices::ChoiceMap, choices::ChoiceMap)
    prev_choice_submap_iterator = get_submaps_shallow(prev_choices)
    prev_choice_value_iterator = get_values_shallow(prev_choices)
    choice_submap_iterator = get_submaps_shallow(choices)
    choice_value_iterator = get_values_shallow(choices)
    choices = DynamicChoiceMap()
    for (key, value) in prev_choice_value_iterator
        key in keys(choice_value_iterator) && continue
        set_value!(choices, key, value)
    end
    for (key, node1) in prev_choice_submap_iterator
        if key in keys(choice_submap_iterator)
            node2 = get_submap(choices, key)
            node = update_recurse_merge(node1, node2)
            set_submap!(choices, key, node)
        else
            set_submap!(choices, key, node1)
        end
    end
    for (key, value) in choice_value_iterator
        set_value!(choices, key, value)
    end
    for (key, node) in filter((k, _) -> !(k in keys(prev_choice_submap_iterator)), choice_submap_iterator)
        set_submap!(choices, key, node)
    end
    return choices
end

Can I expect that every iterator which is returned from get_submaps_shallow and get_values_shallow has keys defined as a method for dispatch?

This method above performs a merging which won't throw an error (unlike the method implemented for Base.merge) - but instead non-destructively merge the second over the first.

And a question for discard:

function update_discard(prev_trace::Trace, choices::ChoiceMap, new_trace::Trace)
    discard = choicemap()
    prev_choices = get_choices(prev_trace)
    for (k, v) in get_submaps_shallow(prev_choices)
        get_submap(get_choices(new_trace), k) isa EmptyChoiceMap && continue
        get_submap(choices, k) isa EmptyChoiceMap && continue
        set_submap!(discard, k, v)
    end
    for (k, v) in get_values_shallow(prev_choices)
        has_value(get_choices(new_trace), k) || continue
        has_value(choices, k) || continue
        set_value!(discard, k, v)
    end
    discard
end

Here I have to check using get_submap and an explicit type check - that seems inefficient? But I don't think ChoiceMap supports has_submap for example.

[femtomc]

@marcoct I'm extracting the WithProbability combinator into another PR. Now, I'm working on AD and then I'll setup a unit test suite.

[georgematheos]

@femtomc @marcoct this looks great!

To clarify--it sounds like you have implemented the instance where it is possible for branches to share addresses, right?

If so, what performance impact does this have in the case where they cannot share addresses? In the case where eg. they are static functions we can be sure don't share addresses, do we still have to do full choicemap scans when switching branches?

[marcoct]

If so, what performance impact does this have in the case where they cannot share addresses? In the case where eg. they are static functions we can be sure don't share addresses, do we still have to do full choicemap scans when switching branches?

@georgematheos I think switching branches for now will always need to involve a full execution of the new branch generative function. I could imagine an optimization pass that is applied to the SML + combinators IR code that replaces the branch combinator call with a call to a generative function implementation that is specially compiled to know about both branches and switch more efficiently, but I think that would be a separate piece of code from this combinator implementation.

Edit: @georgematheos Sorry, after thinking more I think I understand your question now -- you're right that there will be performance cost to doing the choice lookups. I think the answer for this is probably to (in a separate PR) add flag(s) that instruct the combinator to use separate address spaces, or informs the combinator that the address spaces are disjoint.

[femtomc]

@georgematheos it's also possible to optimize this in a more fine-gained automatic way with @generated functions for the two utility methods update_recurse_merge and update_discard (and a similar traversal/construct function in the regenerate interface).

These two functions are responsible for traversing and identifying which addresses need to be constrained. But if we have access to static address information, it should be possible to specialize these on static information. Is this what you might have been thinking about? Do SML functions carry their support as type information (or otherwise allow me to maybe call a function which has cached the support information as dispatch on the unique model or trace type)? I think the answer is yes, but I haven't looked into the internals over there.

I was actually going to reach out to you, because some things I noticed about the implementation of regenerate reminded me of your open PR #279 - for regenerate: I perform a move now which requires that I merge but omit addresses which are in Selection - so I've expressed this by constructing a ChoiceMap with this in mind and calling generate along the new branch, and then there's some fancy Selection logic to compute the correct weights, etc.

For this combinator, the notion of update and regenerate are actually very similar when you switch branches (because there's the possibility that there exists a set of addresses which are not constrained in the new space). So update allows you to switch (and constrain in the new space) while regenerate allows you to switch (and eliminate constraints). It seems that these could be unified in some way.

[femtomc]

@georgematheos also, this combinator superficially reminds me of the Dispatch combinator which you, keorn and myself had discussed a long time ago. I wonder if the implementation would be similar?

[femtomc]

As I thought about this more, I'm a little uncomfortable with providing a specialization which depends upon the modelling language. This doesn't seem to match the philosophy of combinators providing language-agnostic patterns.

It becomes a bit tricky to think about where you would put this functionality. This is not something which you would implement to extend the existing ChoiceMap interface - because it requires static information about the full support, and not just a single execution.

It reminds me of some of the new interfaces I'm designing for my modelling language Jaynes - where I've made a point to derive and associate this information with models (e.g. this static information would be succinctly captured in trace_type for my typed extension of GenerativeFunction - for models where you can derive and characterize this type.

These interfaces AFAIK are not a part of the normal methods you implement now for inheritors of GenerativeFunction - but the core requirement would basically include compile-time computation of address space overlap.

So I've been rambling - but maybe my proposal is to introduce this function as part of the GenerativeFunction interface now - and for DML functions you can't specialize (so you fallback to the current implementation) but SML functions can specialize, and so can other trace type carrying inheritors of GenerativeFunction.

Is there an easier way?

[georgematheos]

@femtomc yes--I think that we could implement Dispatch using Switch. And I have a fork of Gen with lots of changes, including merging update and regenerate into a single method (#282); @marcoct and I still need to discuss more concretely if this will become an official part of Gen, and if so, what the timeline is for that.

[georgematheos]

I think the idea of utilizing type information of Generative functions to specialize is a pretty promising idea. I have not fully thought through it but it seems like for static generative functions, it would be able to help. For things like Map, etc., which have a predictable choicemap structure but varying specific addresses, it's not yet clear to me whether an approach like this could work.

I like the idea you are bringing up about having more static information about the generative function in the trace types; I haven't developed many concrete proposals/ideas beyond what you showed me last time we discussed this, but as I said I think it could be useful for performance, static analysis, etc.

[georgematheos]

In terms of a potential "easier way" to implement Switch which fits within the current GFI and does not require doing a choicemap scan every time we switch branches, one proposal is:
Make the choicemap for the Switch(...) object will be the choicemap of the activated generative function's trace, but nested at trace address i =>, where i is the index of the current branch.

This way, there can never be shared addresses between branches, so when switching branches, we just return the old choicemap as the discard, and generate a new choicemap (from scratch/the update constraints).

I could imagine implementing this version of Switch, then also having a ShallowSwitch which could call Switch internally, but have logic to hide these i => prefixes in the choicemap, and handles the logic of copying parts of the choicemaps over as needed (eg. using update_recurse_merge and update_discard).

The idea here is that this gives users a performance vs readability tradeoff. They can use Switch which slightly more cluttered namespaces when they need better performance, and ShallowSwitch otherwise.

I could imagine (potentially in a future PR) making it so that ShallowSwitch can additionally accept static information about the choicemap types of each generative function it is called on, so it can produce more efficient implementations when that static information is available. If we did this, we could for instance translate if statements in the static DSL into calls to ShallowSwitch which know the specific addresses that can be sampled in each branch. If this is tentatively part of the plan, I am less sure if it is valuable to have factorized into Switch and ShallowSwitch, though.

[marcoct]

@femtomc This looks mostly really good! I have a few cosmetic comments, but I found what I think is a bug in the computation of the discard in update.

[marcoct]

I think it should be possible to use get_selected, which is already implemented, instead. It has the same behavior as regenerate_recurse_merge:

get_selected(get_choices(state.prev_trace), complement(selection))

[marcoct]

It might be better to use Diff instead of UnknownChange here. There might be other Diff types that could be passed that are intermediate between UnknownChange and NoChange (e.g. there is an IntDiff already, which tracks the arithmetic difference between two integers). This method should apply to anything that's not a NoChange.

[marcoct]

I wonder if it would be simpler and more efficient to do:

for (key, node2) in choice_submap_iterator
    if isempty(get_submap(new_choices, key))
        set_submap!(new_choices, key, node2)
    end
end

[femtomc]

@marcoct This one is slightly interesting: the intent here is to directly identify what submap addresses in choices are not in prev_choices. The submaps associated with these addresses can be directly copied over. Addresses which do conflict are treated by the recursive call above these lines.

This suggestion is certainly simpler - is there an argument for efficiency?

[marcoct]

Re efficiency, my thinking is that there might be some overhead in the select and complement operations that the second version doesn't have.

[marcoct]

I'm also dispatching on process! using the diff types. I'm not sure if this current implementation is correct - mostly because I'm not supported Diffed dispatch yet (and I suspect I have to).

@femtomc The Diffed types are only used for boxing Julia values for diff propagation of Julia code via operator overloading. GFI implementers only need to worry about Diff values. So what you have is right.

[marcoct]

A comment here or a docstring for update_recurse_merge that describes the operation this does would be helpful. For example, maybe "Returns choices that are in constraints, merged with all choices in the previous trace that do not have the same address as some choice in the constraints.".

[femtomc]

@marcoct Added a @doc - is this appropriate for formatting documentation in the project? Or should I use a docstring directly?

[marcoct]

I think it's fine to use @doc.

[marcoct]

Docstring or comment would be helpful here, e.g. "Returns choices from previous trace that (i) have an address that does not appear in the new trace, or (ii) have an address that does appear in the constraints (choices)."

[femtomc]

Same as above - added an @doc.

[marcoct]

I found this hard to read. I think that writing out the logic with explicit if-else would be clearer here.

if has_value(get_choices(new_trace), k) && has_value(choices, k)
    set_value!(discard, k, v)
end

But also, I think this should be changed to:

if (!has_value(get_choices(new_trace), k)) || has_value(choices, k)
    set_value!(discard, k, v)
end

[femtomc]

Yes - this is a lot simpler and easier to understand.

[femtomc]

There was actually an error here that you caught.

You should add a choice to the discard if:

1. The choice is in the old trace but not in the new trace.
2. The choice is in the old trace and (the choice is in the new trace and in the constraints).

Previously, I was checking for 2 but I wasn't covering 1.

[marcoct]

I think this requires a recursive call.

Even if the new trace has some choice under this key, and the constraints (choices) also has some choice under this key, there could also be choices in the new trace that were copied directly from the previous trace and not overwritten by the constraints.

[femtomc]

You're absolutely right - this was a mistake in my logic.

[femtomc]

Excellent - will address these tomorrow.

[femtomc]

@marcoct Hi Marco - I wasn't sure what the protocol was for addressing review comments (e.g. if I should show code directly in comments to resolve conversations). I've pushed the necessary changes - so you might examine the changes and determine if they have resolved what you've suggested.

I've also added a test case for update_discard.