AePPL, discrete mixtures and conditionals

Let us consider a family of probability distributions $\left\{{\mathbb{P}}_{\mu }:\mu \in \mathcal{M}\right\}$ where $\mathcal{M}=\left\{1,\dots ,K\right\}$ and $\mathbb{Q}$ a probability distribution on the domain $\mathcal{M}$. A discrete mixture model assumes that the value of a random variable $Y$ can be drawm from one of the ${\mathbb{P}}_{\mu }$: at each step we draw one value $m\in \mathcal{M}$ from $\mathbb{Q}$, and then draw from ${\mathbb{P}}_m$.

We can then write the likelihood of $Y$ conditional on $m$ as:

$$\mathbb{P}(Y\mid z)={\mathbb{P}}_z(Y)$$

Many probabilistic programming languages implement mixture via an ad-hoc Mixture distribution object. In Aesara/AePPL mixtures are expressed via the generative process described above, using the following Aesara constructs:

• Indexing an array with a random variable;
• aesara.tensor.where;
• aesara.ifelse.ifelse.

Indexing an array

We can define mixture by using a random variable with a discrete support (Bernoulli, Categorical, etc.) to index an array of random variables:

import aeppl
import aesara
import aesara.tensor as at
import numpy as np
 
srng = at.random.RandomStream(0)
 
loc = np.array([-1, 0, 1, 2])
N_rv = srng.normal(loc, 1.)
 
p = np.array([0.2, 0.3, 0.1, 0.4])
I_rv = srng.categorical(p)
 
Y_rv = N_rv[I_rv]
 
sample_fn = aesara.function((), Y_rv)
print(sample_fn())
 
logprob, (y_vv, i_vv) = aeppl.joint_logprob(Y_rv, I_rv)
logprob_fn = aesara.function((y_vv, i_vv), logprob)
print(logprob_fn(10, 0))

Using aesara.tensor.where

We can also define mixtures using aesara.tensor.where; the conditional can be based on other random variables:

import aeppl
import aesara
import aesara.tensor as at
 
srng = at.random.RandomStream(0)
 
x_rv = srng.normal(0, 1)
y_rv = srng.cauchy(0, 1)
i_rv = srng.normal(0, 2)
 
Y_rv = at.where(at.ge(i_rv, 1.), x_rv, y_rv)
 
sample_fn = aesara.function((), Y_rv)
print(sample_fn())
 
logprob, (y_vv, i_vv) = aeppl.joint_logprob(y_rv, i_rv)
logprob_fn = aesara.function((y_vv, i_vv), logprob)
print(logprob_fn(10, -1.))

TODO Using aeara.ifelse.ifelse

We also soon be able to use aesara.ifelse to define mixures in the same way we use aesara.tensor.where:

import aeppl
import aesara
import aesata.tensor as at
 
srng = at.random.RandomStream(0)
 
x_rv = srng.normal(0, 1)
y_rv = srng.cauchy(0, 1)
i_rv = srng.bernoulli(0.5)
 
Y_rv = aesara.ifelse(i_rv, x_rv, y_rv)
 
logprob, (y_vv, i_vv) = aeppl.joint_logprob(Y_rv, i_rv)