Implement a general mixture log-probability rewrite

[brandonwillard]

This PR implements log-probabilities for general mixtures (e.g. heterogeneous RandomVariables).

This implementation does not produce a marginalized log-probability.

It covers cases like the following:

import aesara.tensor as at

from aeppl import joint_logprob


X_rv = at.random.normal(0, 1, name="X")
Y_rv = at.random.gamma(0.5, 0.5, name="Y")

p = at.scalar("p")
I_rv = at.random.bernoulli(p, name="I")

M_rv = at.stack([X_rv, Y_rv])[I_rv]

m = M_rv.type()
M_logp = joint_logprob(M_rv, {M_rv: m})

Instead of producing a log-probability graph that represents the computation np.stack([logprob(X, obs), logprob(Y, obs)])[idx], which requires complete evaluation of logprob(X, obs) and logprob(Y, obs) even though idx will only select some of those logprob values from each mixture component, this implementation will compute something like the following:

res[idx == 0] = logprob(X[idx == 0], obs[idx == 0])
res[idx == 1] = logprob(Y[idx == 1], obs[idx == 1])

Currently, it does all the necessary setup and parsing needed to construct the desired log-probability graphs; however, the actual construction steps aren't necessarily finished. I still need to check/confirm a few things and—obviously—write tests.

[brandonwillard]

The multivariate case is blocked by aesara-devs/aesara#507 and aesara-devs/aesara#509.

The latter issue is only a blocker if we want to prevent the evaluation of "unselected" mixture components. We can always produce an array of stacked logprob values evaluated conditional on each mixture component and then index that array with the mixing distribution; however, this is a prohibitively slow approach for large enough mixture models (e.g. almost anything involving time-varying mixture components used with non-trivially-sized datasets).

[brandonwillard]

The two tests in this PR pass locally under the changes in aesara-devs/aesara#516.

[codecov]

Codecov Report

Merging #19 (a74b7ec) into main (b07159b) will increase coverage by 2.27%.
The diff coverage is 93.75%.

@@            Coverage Diff             @@
##             main      #19      +/-   ##
==========================================
+ Coverage   90.52%   92.80%   +2.27%     
==========================================
  Files           5        6       +1     
  Lines         802      917     +115     
  Branches      103      119      +16     
==========================================
+ Hits          726      851     +125     
+ Misses         44       30      -14     
- Partials       32       36       +4     

Impacted Files	Coverage Δ
aeppl/joint_logprob.py	98.00% <83.33%> (+6.33%)	⬆️
aeppl/mixture.py	93.63% <93.63%> (ø)
aeppl/opt.py	82.60% <100.00%> (+22.00%)	⬆️

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update b07159b...a74b7ec. Read the comment docs.

[ricardoV94]

We could have a rewrite phase before this one where we convert constants to Dirac deltas (not needed for this PR)

[ricardoV94]

Wouldn't be safer to fill it with nans?

[brandonwillard]

Safer in what way? It probably shouldn't matter what that initial value is, since the indexing is supposed to cover all values.

[ricardoV94]

Is this working with categorical distributions as well? If so, we could add a test

[brandonwillard]

Is this working with categorical distributions as well? If so, we could add a test

Yes, definitely.