Lift random variables through mixtures when looking for closed-form posteriors

[rlouf]

Copying @brandonwillard's message from #4. If we lift random/measurable variables through mixtures, we can enable some important closed-form posterior opportunities.

For example:

import aesara
import aesara.tensor as at


srng = at.random.RandomStream(4238)

I_rv = srng.bernoulli(0.5, name="I")

Z_1_rv = srng.gamma(10, 100, name="Z_1")
Z_2_rv = srng.gamma(1, 1, name="Z_2")

Z_rv = at.stack([Z_1_rv, Z_2_rv])

# Observation model
Y_rv = srng.poisson(Z_rv[I_rv], name="Y")

Conjugate updates are available between Y_rv and the two Z_*_rv, conditional on the values of I_rv.

The model after lifting should be equivalent to the following:

Z_1_new_rv = srng.poisson(Z_1_rv, name="Z_1_new")
Z_2_new_rv = srng.poisson(Z_2_rv, name="Z_2_new")

# New observation model
Y_new_rv = at.stack([Z_1_new_rv, Z_2_new_rv])
Y_new_rv.name = "Y_new"

The Z_*_new_rv terms are now amenable to the Poisson-gamma conjugate rewrites.