Expand models supported by automatic marginalization

[ricardoV94]

This PR allows more kinds of graphs to be marginalized. Previously, we were limiting it to Elemwise operations to ensure that information across batch dimensions was not mixed between the marginalized and dependent RVs, so as to generate an efficient logp expression that did not grow with number of batch dimensions, but was constant on the domain of the marginalized variables.

We still have the same constraint, but can now analyze it more carefully by propagating information about the batch dimensions of the marginalized RV across the intermediate operations. This allows operations like DimShuffle (transposition, expand_dims, squeeze), Blockwise, Reductions and, rather important, flavors of basic and advanced indexing like the following:

import pymc as pm
import pytensor.tensor as pt

from pymc_experimental import MarginalModel

with MarginalModel() as m:
    state = pm.Categorical("state", p=[0.1, 0.3, 0.6], shape=(4,))
    # Advanced indexing was not supported before
    # The indexed variable could be an RV as well!
    mu = pt.as_tensor([-10, 0, 10])[state]
    sigma = pm.HalfNormal("sigma")
    emission = pm.Normal("emission", mu, sigma, observed=[-9.0, -0.5, 1.0, 11.0])
    
m.marginalize(state)
m.point_logps()
# {'sigma': -0.73, 'emission': -10.52}

This should expand the range of models supported and open room for further expansions.

Some limitations were introduced to simplify the internal logic such as:

1. Marginal RVs can't have explict broadcastable batch dims, but can be expanded explicitly or implicitly downstream. Thus, pm.Bernoulli(..., shape=(5, 1)) is not allowed, but pm.Bernoulli(..., shape=(5,))[:, None] is. This simplifies the logic by only having one kind of possible dim in each axis.
2. Dependent RVs with batch dims beyond those introduced by the Marginal RV, must have such batch dims to the left

TODO:

• Give more informative error for explicit broadcastable dims
• Document internal logic more carefully
• Allow dependent multivariate RVs (it should almost work out of the box, just need some tweaks and tests)
• Test subgrah_dims utility directly
• Test MvNormal k-clusters model that showed up in Discourse some time ago
• Test dependent RVs batch to the left restriction
• Allow a couple more simple Ops

[zaxtax]

What does other mean here in terms of the variable name?

[ricardoV94]

it's the set of model RVs that are direct ancestors to the dependent rvs, excluding the marginal rv

[zaxtax]

It would be nice if there was some intuition for why some programs can be marginalized and some cannot. The distinction is a bit and for some cases it feels like we should be able to.

[ricardoV94]

It would be nice if there was some intuition for why some programs can be marginalized and some cannot. The distinction is a bit and for some cases it feels like we should be able to.

You can marginalize models as long as they don't mix dimensions of the marginal RV (and as long as we can be sure of that).

So if idx is the variable you are marginalizing, a direct dependent variable could use as a parameter idx + idx or idx.T + idx.T, but not idx + idx.T because it mixes distinct dimensions. Similarly, something like sum(idx) is not allowed.

[ricardoV94]

This PR now depends on pymc-devs/pymc#7159

[zaxtax]

Aside from being re-based, are there any blockers to merging this?

[twiecki]

Aside from being re-based, are there any blockers to merging this?

@ricardoV94 just had a baby so probably won't respond. Do you want to rebase and we can just merge for now?

[ricardoV94]

Nah I can still type now and then. PR was not yet done

[zaxtax]

Congrats @ricardoV94 🥳