[GJEP] Dynamic addresses

[femtomc]

This is a reference issue which ties together several other design issues related to "dynamic addresses" in GenJAX, including #556, #473, #239 -- even formalization issues like #235 (which interact with this feature, by virtue of the fact that this feature requires robust masking, described in a moment) and #212.

Quick note:

In the following, I'm using a new convention gen fn construction convention and naming -- (to make clear what language is in use) -- Static below refers to what I'm called in other issues / thread as GenJAX's Builtin generative function language. Static refers to the fact that "JAX transforms are active in the interface implementations".

Dynamic addresses

Fundamentally, this issue is about the following type of generative function:

@gen(genjax.Static)
def model(t: IntArray):
    x = genjax.tfp_normal(0.0, 1.0) @ ("x", t)
    return x

Seemingly innocuous -- except it's not: it uses t - which can be a dynamic, JAX traced, integer array as part of the address.

On first design, addresses of this form were explicitly not supported in GenJAX's Static language -- the interfaces were implemented using Trie-like types, whose keys had to be hashable. In first design, address components were constrained to be statically hashable things - including strings, and Python literal ints.

However, @derifatives quickly noted that expressing SMC update steps over sequences of target distributions expressed via genjax.UnfoldCombinator models is inconvenient with Static language programs whose addresses cannot be dynamic. Why? Well, you'd really like to write a proposal (like the program above):

@gen(genjax.Static)
def proposal(t: IntArray):
    latent= genjax.tfp_normal(0.0, 1.0) @ ("above_an_unfold", t, "latent")

And then, inside SMC -- write a tight jax.lax.scan -- which provides the current SMC target index t to the proposal, to produce a choice map which targets that address.

If the choice maps from the Static language are disallowed from containing dynamic integer addresses, you cannot construct such a jax.lax.scan - based SMC program - because jax.lax.scan always provides dynamic values to the body function (meaning, you can't use t in the address, because it's dynamic -- and it will always be dynamic -- inside of jax.lax.scan).

There's several "shallow" ways to fix this problem. One way would be to force users who use SMC to provide a proposal which doesn't actively include the current SMC target step t in their address space -- and then massage their resulting choice maps to work correctly at the current time step. By hoisting t out of the proposal entirely, we don't have to worry about dynamic addresses. On the other hand, you often do want to know what t is so e.g. you can write proposals that look at the current observations. More pressingly, removing t from user control restricts the user from writing more expressive proposals. As part of Gen's philosophy, we want to err on side of supporting expressivity as much as possible (and, within reason, for our implementation's limitations).

Dynamic addresses are infectious

Dynamism in JAX is a tricky property. Where dynamism occurs, code cannot make a decision at trace time, and instead must be conservative. Dynamism immediately applies that we have to push a branching decision into runtime somewhere. We have control over the somewhere, but we cannot eliminate the branching.

As mentioned in other issues, the ChoiceMap interfaces like get_submap and has_submap must contend with dynamism in certain settings. One example is in genjax.SwitchCombinator - where "what branch generative function was called" is a dynamic quantity. To handle dynamism in his context, GenJAX has a type called Mask (#235, https://probcomp.github.io/genjax/genjax/library/core/functional_types.html?h=mask#option-types-the-masking-system): Mask wraps a dynamic piece of data with a dynamic Bool tag -- it's just like a functional Maybe or Option type.

If one invokes SwitchChoiceMap.get_submap -- and the address is not valid for the branch which was actually taken, we need some way to denote that behavior. We always return some value - but because we can't know what branch was actually active until runtime, we need to tag that value with a piece of information to say whether or not it is valid or invalid at runtime. We do so using Mask - which always wraps a value, but the "validity" of the value depends on the tag.

But surprise surprise! Dynamism is part of dynamic addresses -- if one constructs a choice map which supports dynamic addresses, and a user asks for DynamicChoiceMap.get_submap with a runtime integer array, whether or not that integer array is actually in the choice map is dynamic. Thus, dynamic addresses will likely involve at the very least, some form of masking.

Let's explore what else is required below.

Dynamic addresses and update

Furthermore, we can quickly consider an extreme manifestation of the "infectiousness" of dynamism -- let's say you write:

@gen(genjax.Static)
def model(t: IntArray):
    x = genjax.tfp_normal(0.0, 1.0) @ ("x", t)
    return x

and generate a trace from this generative function.

If you attempt to update this trace in JAX traced code, and you change t -- what happens?

Well, Gen allows you to change the measure, including the addresses by changing the arguments. But now, the address which the update interface sees might imply that a subtrace (for the callee, here a Distribution generative function tfp_normal) either exists or does not exists.

That means that the update interface must handle dynamism by introducing branching -- because whether or not a subtrace lives at that address is now a runtime decision.

I cover an example like this in more detail in my replies to this issue thread (below).

Solutions: new dynamic choice map type and masking

To take a step towards the solution space, I believe that a design solution will consist of a combination of features:

• (Masking) The masking system -- already implemented -- must be stress tested and made robust, as fully supporting dynamic addresses in GenJAX's StaticLanguage will involve heavy use of masking.
• (Specialization) When dynamic addresses are not active, the user should not incur the cost of dynamic addresses. This can be solved at trace time by specializing choice map types to "concrete" choice maps (read: choice maps without any dynamism).
• (Dynamic choice map) We must introduce a new type of choice map which can natively handle dynamic addresses via careful usage of Pytree -- to separate dynamic and static data.
• (Interface code to handle dynamic choice maps) This last point may be the most important, and it is reflected in issues like [Defactor] Support Builtin proposers for Vector combinators #212, and [Defactor] Choice map conversion and generic fallbacks #556. The new type of dynamic choice map, in a very concrete sense, is the most general type of choice map. It can support hierarchical addresses of any structure. By virtue of this fact, it's worth supporting code paths for all existing languages to accept dynamic choice map. This provides a common conversion target for generative function languages which utilize specialized choice maps (e.g. interoperability between the choice map types of different languages, which is required by the Gen's semantics -- which doesn't specialized on the data representation of a choice map, but does require interoperability of the address space of a choice map (no matter the representation)). Without further experience, my prediction is that the cost of conversion which be a tracing time cost -- and utilizing dynamic choice map will incur runtime penalties due to branching. But fundamentally, strong code paths for this functionality supports "slow, but stable and compositional" interface implementations

The criteria for completeness involves the following testing axes:

• (No regressions) Introducing dynamic addresses as a feature should lead to no performance regressions or unintended errors for existing code that does not utilize dynamic addresses. This is a test of specialization above.
• (Compositionality - robustness tests for simple models) By virtue of the fact that we are introducing a new feature, we must stress test the compositional usage of this feature across a variety of scenarios.
• (Compositional - SMC) By virtue of the fact that we are introducing a new feature to solve a set of existing problems, we must provide a new set of tests to illustrate usage of the feature (without errors) to solve those problems.

[femtomc]

As mentioned in the section on update above - there are pathological settings when using dynamic addresses that we need to control for using either (a) runtime checks or (b) explicit restrictions on how dynamic addresses may be used within the Static modeling language.

Consider the following model:

@gen(Static)
def model(t):
    x = normal(0.0, 1.0) @ ("x", t)
    _ = normal(0.0, 1.0) @ ("x", t + 1)

Now, this is a valid a model -- the addresses will always be different (so it doesn't violate the no duplicate address restrictions).

However, when we consider update, we now need to be very careful:

model.update(key, prev, new_constraints, (Diff(t, NoChange), ))

The way that update is implemented for the Static language means that we must fetch a callee's trace from the previous (StaticLanguageTrace). When we hit the address ("x", t) -- a problem arises: we do not know the value of t at JAX tracing time -- which means we cannot distinguish ("x", t) from ("x", t + 1). Therefore, we have to be conservative -- we must fetch both subtraces.

Now we have another implementation choice to make: we can insert a jax.lax.cond only if the return Pytree types of the update call to the callee's traces (both of them, because we have to deal with both) are the same.

We've encountered a "switch combinator"-like scenario inside our implementation of StaticLanguageGenerativeFunction.update -- and this scenario has arisen just because of the properties of dynamism.

Now, of course, just like SwitchCombinator -- only one of the update invocations is actually valid at runtime. So one solution would be to make a type that is like SwitchTrace (maybe a generic version, which is not so tied to SwitchCombinator) and shove the two updated traces inside (with an array integer tag denoting which is valid).

The larger problem is that the number of calls is quadratic in the number of possibly conflated dynamic addresses. So e.g. if you have 4 of these addresses in your model code, each time you update -- for each of those addresses, you're going to to have to run an update for all of them.

[femtomc]

For the above issue, one possible solution I came up with — take advantage of the fact that the traversal order of addresses in GenJAX is deterministic.

The pathological case in that solution is if a change to an argument changes the address support — but I believe we can handle that with a simple cond.

[femtomc]

Here's the interesting case to consider in the above solution:

Again, consider the following model:

@gen(Static)
def model(t):
    x = normal(0.0, 1.0) @ ("x", t)
    _ = normal(1.0, 1.0) @ ("x", t + 1)

If we generate a trace with t = 0, it has 2 addresses (("x", 0) and ("x", 1)). Suppose we update this trace by changing the t argument, but providing no new constraints.

tr.update(key, empty_choice(), (Diff(1, UnknownChange), ))

Now, at the first address, we have ("x", 1) -- but the old address by traversal order in the trace is ("x", 0). However, ("x", 1) is in the trace (at the second visited address). This means we have to get the subtrace at that address, and call update on it.

So it does not suffice to just grab the subtrace at the index in the list of addresses by traversal order -- because changes to the arguments can affect what subtrace we want to grab. (e.g. -- because we've shifted the address -- we need to grab the second callee, not the first -- but making that decision might by dynamic!)

[femtomc]

Based on my reasoning, here are the degrees of freedom we have available to us for this feature:

• The simulate, importance, and assess methods suffer small changes -- mostly in the data types produced by the language (e.g. now we have a new type of DynamicHierarchicalChoiceMap -- which natively works with dynamic addresses). In the worst case -- when a user provides a dynamic choice map to importance -- and now dynamism creeps into "what address do I need to get?" -- the performance hit is due to a jax.lax.cond at the leaves. There are degrees to the hit -- e.g. "fully dynamic" DynamicHierarchicalChoiceMap is quite bad -- but a user can specify a partially static, partially dynamic choice map like a (HierarchicalChoiceMap with an IndexedChoiceMap -- where the dynamism is over the index dimension) which induces a significantly reduced penalty. As part of the optimization of this language, we might be able to handle conversion (when appropriate) automatically.

• The update interface has several cases.

  • If arguments which affect the address support do not change, we can simply use the traversal order (known statically) to fetch the correct trace -- compared to fetching the trace using the address itself, which (because of dynamism), implies that we can't resolve which trace exactly statically. Using the traversal order, we can resolve this fact statically.
  • If an argument change affects the address support -- now the traversal order may not index correctly into the subtrace we need to grab. If the change to the argument is statically known (e.g. say StaticIntDiff) -- we might be able to shift the traversal order statically -- by defining the correct Diff propagation rules in the incremental interpreter.
  • Finally, as the most general fallback -- any dynamic addresses which are active might require that we merge the trace choices with constraints, and then importance at the new address -- instead of trying to wrangle what it means to update "all traces in the set of possible traces at this dynamic address". Then, instead of trying to do the dynamic song and dance with "what update code we need to call" -- importance just works, and we manage the score math via the discard after the fact, but we lose out on incremental computation in the callee.

In general, I think the trade offs are acceptable. Model changes via dynamic changes to arguments which affect dynamic addresses are the "most exotic" form of dynamism in GenJAX -- and it seems fair to me that the implementation is the least performant around this feature.

[femtomc]

Closing this as "won't address". After internal discussion, I believe that the path forward to support dynamic addresses is to expose them via specific inference functionality (like trace translators), and not to try to push them into the core Static language (making the language harder to guarantee / ensure correctness and bug freedom).