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.

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]

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 🥳

[ricardoV94]

Failing CI is due to weird stuff between PyMC and PyTensor: pymc-devs/pytensor#1009

Hopefully it will be sorted out soon. Tests pass locally on my machine.

Open for review!

[zaxtax]

This looks really good. It's a little annoying that the refactor is mixed in with expansion of support for models, but looks pretty good.

[ricardoV94]

It's a little annoying that the refactor is mixed in with expansion of support for models, but looks pretty good

It's in separate commits

[zaxtax]

It's a little annoying that the refactor is mixed in with expansion of support for models, but looks pretty good

It's in separate commits

Yea, I reviewed it commit by commit. I saw the bulk of the work was in the first one.

[zaxtax]

We should probably bump the minimum version required to 5.17

[ricardoV94]

We should probably bump the minimum version required to 5.17

Done

[zaxtax]

LGTM!