How can I create a class using pm.Model's class like keras.Model or torch.nn.Module option?

I was wondering how to create a class using pm.Model’s class. Since I am coming from Torch and Tensorflow based models, I like to know how to create a class inheritance based on pymc.Model’s class like the keras.Model or torch.nn.Module option. The documentation for the Model’s class is without a clear example, unfortunately. Is there any example post on creating a class using the pm.Model’s class and how to use it?

Link: [pymc.Model](pymc.Model — PyMC 4.4.0 documentation)


The API for PyMC is a bit different from Torch. You don’t ever need to subclass pm.Model to create a model. Instead, it is used as a context manager, as follows:

basic_model = pm.Model()

with basic_model:

    # Priors for unknown model parameters
    alpha = pm.Normal("alpha", mu=0, sigma=10)
    beta = pm.Normal("beta", mu=0, sigma=10, shape=2)
    sigma = pm.HalfNormal("sigma", sigma=1)

    # Expected value of outcome
    mu = alpha + beta[0] * X1 + beta[1] * X2

    # Likelihood (sampling distribution) of observations
    Y_obs = pm.Normal("Y_obs", mu=mu, sigma=sigma, observed=Y)

This example comes from the introduction and overview example notebook in the examples gallery, which I highly recommend checking out.

If you wanted to write a class wrapper around a PyMC model, perhaps to get some additional functionality (like a .fit method), you wouldn’t subclass pm.Model. Instead just write a generic class with a pm.Model as a member variable and proceed as normal.