Any way to "vmap" or "plate" in pymc

[justindomke]

Hello all,

I was wondering, does pymc support any construct similar to jax.vmap? Consider this simple model, in which I've used vectorized distributions:

with pymc.Model():
    z = pm.Normal("z", mu=[0,1], sigma=[2,3])
    x = pm.Normal("x", mu=z, sigma=[4,5])

I was hoping for some kind of function like pymc.vmap that would allow me to do something like this instead:

def get_x_y_pair(a,b,c):
    z = pymc.Normal("z",a,b)
    x = pymc.Normal("x",z,c)
    return z,x

with pymc.Model():
    z, x = vmap(get_x_y_pair)([0,1],[2,3],[4,5])

Of course, for this simple example the original vectorized syntax is much clearer. However, when things get complicated, manual vectorization can get quite tricky, so I was hoping for something like this.

I've noticed that pytensor and aesara now have scan, which I guess is sort of more general than vmap, though I haven't been able to figure out if scan can be used directly when defining pymc models, or it's more something that should be used internally.

[ricardoV94]

Hi, the issue is not the graph vectorization, but how PyMC tries to register variables in a model context. That's the same reason you can't put PyMC distributions inside a Scan function. Instead you need to first convert a PyMC model to a PyTensor graph, operate on that graph and then convert back to a PyMC model.

We have been adding some functionality that allows you to represent a PyMC model a pure PyTensor graph for some time, we use that for pm.do and pm.observe.

I explored some vectorization model transformations sometime ago here: https://gist.github.com/ricardoV94/99c53fbb8b2e9a68e1b2c6c4d761eaf4

[ricardoV94]

No, that function is user-facing. But we aren't doing a great job of listing the API in the official docs if that's what you mean. There is also pytensor.tensor.vectorize that is more user friendly (same API as numpy). But for this application to PyMC models I think vectorize_graph is easier to work with.

[ricardoV94]

Regardless of which one you use, the trick is you have to represent PyMC models in that PyTensor fromat given by fgraph_from_model and work on that representation for the vectorization.

[justindomke]

Thanks again. I was just trying to clarify that there are no other resources available to help someone who might want to do this other than the gist you provided, and perhaps reading the source code for pymc and/or pytensor. Is there documentation on either of those functions, or on what a FunctionGraph is and how it's related to a model, etc.?

[ricardoV94]

The pytensor functions are documented (at least in docstrings). The more functional approach to modifying PyMC models is not documented at all because it is still work in progress. Other explorations from developers:

• Implement some predictive model transformations pymc-experimental#141
• Variationally Informed Parameterization pymc-experimental#276

Would be great to start getting some user-facing documentation though!

[justindomke]

OK, so just to clarify, I couldn't find any documentation of the API in any of these places:

https://pytensor.readthedocs.io/en/latest/library/graph/index.html
https://pytensor.readthedocs.io/en/latest/library/graph/graph.html
https://pytensor.readthedocs.io/en/latest/library/tensor/index.html
https://pytensor.readthedocs.io/en/latest/library/tensor/basic.html
https://pytensor.readthedocs.io/en/latest/library/tensor/math_opt.html
https://pytensor.readthedocs.io/en/latest/library/tensor/basic_opt.html

Is there somewhere else I should be looking?

Edit: I did find some documentation in the code, though! https://github.com/pymc-devs/pytensor/blob/e8693bdbebca0757ab11353f121eed0c9b3acf66/pytensor/graph/replace.py