as_model

pymc_experimental.as_model(*model_args, **model_kwargs)

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.

Examples

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")
    pm.sample()

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

m = basic_model()
pm.sample(model=m)

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

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