Add utility to convert Model to and from FunctionGraph

[ricardoV94]

import pymc as pm
from pytensor import dprint
from pymc_experimental.utils.model_fgraph import fgraph_from_model, model_from_fgraph

with pm.Model (coords={"trial": range(3)}) as m:
    y_obs = pm.MutableData("y_obs", [1., 2., 3.], dims=("trial",))
    mu = pm.Normal("mu")
    sigma = pm.HalfNormal("sigma")
    y = pm.Normal("y", mu, sigma, observed=y_obs, dims=("trial",))

pm.model_to_graphviz(m)

fg = fgraph_from_model(m)
dprint(fg)

... # Manipulate model

new_m = model_from_fgraph(fg)
pm.model_to_graphviz(new_m)

[lucianopaz]

@ricardoV94, thanks for being able to kick this off. I have a bunch of comments, but I feel that most of the design decisions should be discussed elsewhere. I've opened this discussion to track the design stuff. My main critique is that the new Op that you introduced will make it harder to apply rewrites, because they will have to weigh what to do about the FreeRV or ObservedRV or Deterministic wrapping Op. I think that we can continue that discussion in the thread that I linked.

[ricardoV94]

Ready for review!

[lucianopaz]

This looks really great @ricardoV94. I’ve got two comments:

1. This is a small nit, but could you add a todo list with the things that still aren’t covered by the transformation to fgraph. At the moment I can think of the initval, mutable dimensions, but there must be some more.
2. I can’t wrap my head around the special care you take with deterministics. It looks like the rebuilt model might have a different structure of the deterministics than the original one? At least the test that you wrote with y_ and y__ make it seem that way to me

[ricardoV94]

This is a small nit, but could you add a todo list with the things that still aren’t covered by the transformation to fgraph. At the moment I can think of the initval, mutable dimensions, but there must be some more.

Mutable dimensions should be working via the NamedVars.

There's a test here:

pymc-experimental/pymc_experimental/tests/utils/test_model_fgraph.py

Line 69 in 5ce8973

def test_data():

The only missing cases I could think were initvals and nested models, both of which raise explicitly in the conversion function.

I can’t wrap my head around the special care you take with deterministics. It looks like the rebuilt model might have a different structure of the deterministics than the original one? At least the test that you wrote with y_ and y__ make it seem that way to me

Yes the rebuilt model will never have the deterministics "floating outside" of the graph (unless they are a direct view on another model variable), so they show up nicely in graphviz.

This may be different than the original user model, but I don't think it's something we have to worry about preserving.

On the other hand, the fgraph model will never have deterministics in the middle of the graph, they will always be "floating outside". This way they shouldn't interfere with rewrites.

[lucianopaz]

I think this looks well enough to merge. Let’s try it out on the actual applications to see if we need to adapt anything