Add support for closed-form posteriors

[brandonwillard]

We need to copy over the kanren relations in symbolic_pymc so that we can sample directly from closed-form posteriors—when possible. See the basic example here.

Within the general process of constructing a sampler for a model graph (see #3), we could apply these rewrites as well, and—in some cases—simply return a graph of the joint posterior, instead of a sampler graph.

[brandonwillard]

FYI: 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.

[brandonwillard]

@ricardoV94, @rlouf, @zoj613, we should try to get this example working next. It's something that could be set up without too much effort and makes for a great combination of all our efforts.

[brandonwillard]

We should split this off into a bunch of sub-issues for each (group of) closed-form posteriors we want to implement.