pymc_experimental.as_model(*model_args, **model_kwargs)[source]#

Decorator to provide context to PyMC models declared in a function. This removes all need to think about context managers and lets you separate creating a generative model from using the model. Additionally, a coords argument is added to the function so coords can be changed during function invocation

Adapted from Rob Zinkov’s blog post and inspired by the sampled decorator for PyMC3.


import pymc as pm
import pymc_experimental as pmx

# The following are equivalent

# standard PyMC API with context manager
with pm.Model(coords={"obs": ["a", "b"]}) as model:
    x = pm.Normal("x", 0., 1., dims="obs")

# functional API using decorator
@pmx.as_model(coords={"obs": ["a", "b"]})
def basic_model():
    pm.Normal("x", 0., 1., dims="obs")

m = basic_model()

# alternative way to use functional API
def basic_model():
    pm.Normal("x", 0., 1., dims="obs")

m = basic_model(coords={"obs": ["a", "b"]})