This started out as an internal discussion some months ago. Since@ricardoV94 has opened#111, I thought that it would be best to condense all of the discussion here and open it up to everyone that is interested.

Goal

PyMC provides a way to define a generative model (and also non-generative models through the use ofPotentials) and then gives access to automatic ways of drawing samples from the prior, posterior and posterior predictive distributions. Since random variables are now aTensor (at some point viatheano, thenaesara and nowpytensor), we can leverage the computational backend to do rewrites of models. I'll list a few cases that are relevant use cases of model rewrites:

• It allows us to implement ado operator (discussedhere). Where we say, "replace some random variable with a given value"
• It allows us to implement anobserve operator. This would allow us to first define a model, and then say "this rv should have these observed values". Maybedo andobserve are equivalent because I don't know enough about do calculus, but from my naive point of view,observe adds a logp term conditioned on the observed values, anddo simply sets the values but ignores the logp.
• We can define the model in some form and then automatically get different variants. For example, get a model that marginalizes out a variable, exploits a conjugacy relation or changes a parametrization.
• We can also address cases where we want different forms for observed and unobserved variables easily:Implement DiffTransform for RandomWalk distributions pymc#6098
• More generally we can do arbitrary functions that take models as inputs and return new models as outputs. For instance, we can use this for GP predictions, where we replace the GP prior by the GPconditional so that all variables and deterministics downstream now depend on the conditional when you dosample_posterior_predictive. Before, users had to recreate their deterministics, on top of the conditional manually.

PyTensor (andaesara before that) enable ways to rewrite the computational graph. The piece that is missing is to connect PyMCModels with the entities that are used for applying rewrites:FunctionGraphs.

What doPyMC models store?

PyMC models work as bookkeepers of a few things:

• Free random variables
• Observed random variables
• Deterministics
• Potentials
• Dimension names
• Dimension coordinate values or lengths
• Model name to prepend to newly added entities
• A compile configuration (unless this has been deprecated at some point and I'm not aware of it)
• The parent model context if the current model is nested within another one
• Stores a mapping between variables and their transformed counterparts in unconstrained space (used for running inference, but not forward samples)

Many of the above entities are simpleTensorVariables and mappings between them. This means that we could very plainly use aFunctionGraph that takes as outputs the random variables, deterministics and potentials as outputs, and we'd get thePyMC model's induced function graph. We will need to store the rest of the information somewhere to be able to make the leap back into aModel instance.

How to target rewrites?

We need to choose how to track some of the meta information. In particular, how to track whichTensorVariable is a free random variable, an observed random variable, a deterministic or a potential. I see two alternatives

• Store this information in thetag of theTensor
• Wrap theTensor using a newOp (the approach taken inAdd utility to convert Model to and from FunctionGraph #111)

The benefit of the first is that rewrites don't need to reason about newly inventedOps since they can work with whatever was in the computational graph. The benefit of the second approach is that the newOps can include extra information that conditions the shapes and value variables of the resulting RVs (e.g.#111 includes the dimension information in theOps directly).

There is still a lot of extra information that we need to carry around with us: the mappings, the configuration and scopes. All of these could potentially be included throughFeatures that are appended to theFunctionGraph.