mcx
mcx copied to clipboard
WIP: Implement a `denormalize` custom Jaxpr operator simplifying MCX logpdfs
Overview
This PR is implementing a generic denormalize
decorator which removes normalizing constants in a logpdf. Per call to contributions in #65.
Implementation
The current implementation is a two passes algorithm:
- Forward pass to find all constant variables in the Jaxpr graph;
- Backward pass to find all simplifying assignment, but skipping
add
andsub
operations where one of the input is a constant;
Once the latter simplifying mapping is found, the rest of decorator code is just a simple execution pass on the Jaxpr, skipping the operations where a simplifying mapping exists.
Limitations
Even though we try to have a fairly generic implementation, some simplifications are not supported at the moment. For instance, we do not propagate constant simplification in concat
or mul
operations. These cases could be supported in the future, if it happens to be a performance bottleneck in MCX.
@rlouf As we discussed on Slack, there is quite a bit of additional complexity to add to this PR to handle properly the support select
condition appearing in lot of distributions logpdf.
I'll start with a fairly dummy implementation, getting it working, and I think then we can iterate on it to make it less naive and using more properly symbolic programming concept (I started looking at Oryx codebase on that).
That sounds like a very good plan to me! I'll have a closer look too when my big PR is merged.