pyro
pyro copied to clipboard
Add a compatibility shim for using contrib.funsor with existing models
This PR makes a few small changes to Pyro's EnumMessenger, MarkovMessenger, EnumAllocator and DimAllocator which allow the use of existing Pyro models and primitives with pyro.contrib.funsor (especially @ordabayevy's new inference algorithms like pyro.contrib.funsor.infer.TraceMarkovEnum_ELBO) without requiring the use of pyroapi.
The changes mostly serve to guarantee that the global state in pyro.contrib.funsor.runtime is synchronized with Pyro's existing dimension management state, primarily pyro.poutine.runtime._ENUM_ALLOCATOR and pyro.poutine.runtime._DIM_ALLOCATOR.