PyTensor backend suggestion

[ricardoV94]

Hi devs,

Me and @jessegrabowski spent an afternoon (I added another on my own) rewriting part of the core library to be built on top of PyTensor, the library that powers the perhaps more famous PyMC PPL, two libraries we maintain.

The fork with non-cleaned up commits is here: https://github.com/ricardoV94/probabilit
Diff: main...ricardoV94:probabilit:main

This was done mostly for the fun of it (and to create new forms of documentation for PyTensor, as we recorded it for later publishing).

We actually think PyTensor is a great fit for a library like yours. It is lazy by default, emulates most of NumPy API, has clear RV signature semantics, and it offers graph walking and rewriting abilities that make it sort of trivial to introduce the QMC draws / correlation logic you are doing in your sample routines.

It also allows you to compile the final function to C/Numba/JAX (and more experimental MLX/PyTorch), which can be great for parallelization over QMC draws / running on GPU.

What we changed

We introduced some major changes that we felt worked better with PyTensor:

1. sample is a function and not a method
   1. It takes the [nodes] which are defined in terms of PyTensor Operations (including RVs for the distributions). It walks the graph and replaces each RVs by the respective quantile function evaluated on a slice of the QMC draws. Finally we vectorize this graph over the size dimension of the QMC draws. This is compile into a PyTensor function with signature qmc_draws: np.ndarray[size, d] -> Sequence[node_draws: np.ndarray[size, *node_shape]]`.
2. You have to select which nodes you want to keep samples for when you call sample.
3. We don't store intermediate draws in the nodes themselves
4. correlate is a function with signature (nodes, vars_to_correlate, corr_matrix, method), that returns a tuple[new_nodes, correlated_vars]. 1. It introduces a dummy Op that is materialized when we vectorize the original graph with the QMC draws
   1. You could actually correlate any (independent) nodes you wanted, not only RVs. We could add check that you're not correlating a function of and the RV itself.

We didn't do a complete refactor, but all tests in test_modeling.py pass except for the the EmpiricalDistribution which we didn't implement.

Demo

Here is your first example from the readme:

from probabilit.sampling import sample
from probabilit.distributions import Distribution

male_height = Distribution("norm", loc=176, scale=7.1)
female_height = Distribution("norm", loc=162.5, scale=7.1)
statistic = (male_height > female_height)
samples = sample(statistic, 999, random_state=0)
float(samples.mean())
# 0.8988988995552063

If you have JAX installed you can also use it

# Could have a more user friendly API, like `backend="JAX"`
samples = sample(statistic, 999, random_state=0, compile_kwargs=dict(mode="JAX"))
float(samples.mean())
# 0.8988988995552063

built-in PyTensor graph inspection

statistic.dprint()
"""
Gt [id A]
 ├─ normal_rv{"(),()->()"}.1 [id B]
 │  ├─ RNG(<Generator(PCG64) at 0x7F89D1F3F300>) [id C]
 │  ├─ NoneConst{None} [id D]
 │  ├─ 176 [id E]
 │  └─ 7.1 [id F]
 └─ normal_rv{"(),()->()"}.1 [id G]
    ├─ RNG(<Generator(PCG64) at 0x7F89D1F3F680>) [id H]
    ├─ NoneConst{None} [id D]
    ├─ 162.5 [id I]
    └─ 7.1 [id J]
"""

If you're a dev you probably want to see the kind of compiled sample function were generating

samples = sample(statistic, 1, compile_kwargs=dict(mode="JAX"), dprint_sample_fn=True)
"""
Gt [id A] <Vector(bool, shape=(?,))> 15
 ├─ Add [id B] <Vector(float64, shape=(?,))> 14
 │  ├─ [176.] [id C] <Vector(float64, shape=(1,))>
 │  └─ Mul [id D] <Vector(float64, shape=(?,))> 13
 │     ├─ [10.04091612] [id E] <Vector(float64, shape=(1,))>
 │     └─ Erfinv [id F] <Vector(float64, shape=(?,))> 12
 │        └─ Add [id G] <Vector(float64, shape=(?,))> 11
 │           ├─ [-1.] [id H] <Vector(float64, shape=(1,))>
 │           └─ Mul [id I] <Vector(float64, shape=(?,))> 10
 │              ├─ [2.] [id J] <Vector(float64, shape=(1,))>
 │              └─ Subtensor{:, i} [id K] <Vector(float64, shape=(?,))> v={0: [0]} 9
 │                 ├─ qmc_samples [id L] <Matrix(float64, shape=(?, 2))>
 │                 └─ 1 [id M] <uint8>
 └─ Add [id N] <Vector(float64, shape=(?,))> 8
    ├─ [162.5] [id O] <Vector(float64, shape=(1,))>
    └─ Mul [id P] <Vector(float64, shape=(?,))> 7
       ├─ [10.04091612] [id E] <Vector(float64, shape=(1,))>
       └─ Erfinv [id Q] <Vector(float64, shape=(?,))> 6
          └─ Add [id R] <Vector(float64, shape=(?,))> 5
             ├─ [-1.] [id H] <Vector(float64, shape=(1,))>
             └─ Mul [id S] <Vector(float64, shape=(?,))> 4
                ├─ [2.] [id J] <Vector(float64, shape=(1,))>
                └─ Reshape{1} [id T] <Vector(float64, shape=(?,))> v={0: [0]} 3
                   ├─ Subtensor{:, :stop} [id U] <Matrix(float64, shape=(?, ?))> v={0: [0]} 2
                   │  ├─ qmc_samples [id L] <Matrix(float64, shape=(?, 2))>
                   │  └─ 1 [id V] <int64>
                   └─ JAXShapeTuple [id W] <Vector(int64, shape=(1,))> 1
                      └─ Shape_i{0} [id X] <Scalar(int64, shape=())> 0
                         └─ qmc_samples [id L] <Matrix(float64, shape=(?, 2))>
"""

You can work with multidimensional graphs / operations without extra complexity in the codebase.

# Take a draw from the (male, female) populations
height_pop = Distribution("norm", loc=[176, 162.5], scale=1.1)
# Take 4 individuals draws around each of the 2 population draws
height_individuals = Distribution("norm", loc=height_pop, scale=[7.1, 7.05], size=(4, 2))
# You can use any PyTensor operation/method. Here we just transpose
draws = sample(height_individuals.T, 100, method="lhs")
draws.mean(axis=0)
# array([[176.03177248, 175.92640186, 175.99058435, 176.00309828],
#        [162.48392959, 162.44472868, 162.53143526, 162.51765671]])

You can also index, reshape, concatenate and all that sort of stuff for free. I suspect this renders the MarginalDistribution unnecessary.

It's also easy to wrap Ops in PyTensor if you wanted to keep using scipy for distributions not implemented in PyTensor. This forces PyTensor functions to run in Python (unless you provide C/numba/jax impls) but is no worse than before performance wise.

Offer

If you are interested we can easily push this to the finish line and open a PR

Either way, thanks for sharing the library and keep the good work :)

[tommyod]

thanks for sharing the library and keep the good work

Thank you @ricardoV94 and @jessegrabowski. I'm flattered that you both looked at this little library! Although I have used pyMC a little bit a few years back, I have not looked at PyTensor before. I will have to read up on PyTensor and check internally in Equinor if this is a direction we want to go in (since the main purpose of probabilit is to be used in the company). I'm pretty sure it's a technical improvement, but our use cases are quite simple so I don't want to over-complicate it either. Needless to say this is very interesting. If you want to open a PR that would be great, but I cannot say for sure if we would merge it or not just yet.

Again: thank you so much! I'm excited to have a look at PyTensor and your fork in detail!

[tommyod]

Hi again @ricardoV94 ! After thinking a bit and gathering some opinions we've decided that we won't pursue a PyTensor backend. At least not at this time. The reason is that we must prioritize stability (we want the code to run in 10 years with as little effort as possible) and easy of entry/maintenance (many people like myself know numpy/scipy pretty well, but are unfamiliar with PyTensor).

If this was my own project I'd probably use PyTensor, knowing what I know now. It's very interesting and elegant. But the "dumb" scipy/numpy approach might be better for our purposes, which are pretty basic to be honest :) At least that is my assessment right now.

Thanks again! By the way: do you mind if I push a copy of your branch to this repo? That way I can preserve a copy of your work here too. In the future maybe PyTensor will be the way to go for us.

[ricardoV94]

Thanks for the careful consideration. It all sounds pretty reasonable to me.

Feel free to keep a copy/ adapt any of the code as you wish!