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
.