Blackjax On Aesara

import aesara.tensor as at
 
srng = at.random.RandomStream(0)
 
sigma_rv = srng.normal(1.)
mu_rv = srng.normal(0, 1)
Y_rv = srng.normal(mu_rv, sigma_rv)

Sample prior

Sampling from the prior predictive distribution is a useful tool for debugging. One thus needs this function to be:

• fast
• Easy to use
• Easy to customize; is it easy to change the value of a parameter.

import aesara
from aesara.graph.basic import io_toposort
from aesara.tensor.random.op import RandomVariable
import jax
 
rng_key = jax.random.PRNGKey(3)
 
def count_model_rvs(rv_out):
    """Count the number of `RandomVariable` in a model"""
    return len([node for node in io_toposort([], [rv_out]) if isinstance(node.op, RandomVariable)])
 
def split_to_tuple(rng_key, num):
    keys = jax.numpy.split(jax.random.split(rng_key, num), num)
    return tuple([key.squeeze() for key in keys])
 
def prior_sample(rng_key, num_samples, rv_out):
    """Return prior predictive samples"""
    prior_fn = aesara.function([], Y_rv, mode="JAX").vm.jit_fn
    num_rvs = count_model_rvs(Y_rv)
 
    def take_one_sample(rng_key):
        keys = [{"jax_state": key} for key in split_to_tuple(rng_key, num_rvs)]
        return prior_fn(*keys)[0]
 
    return jax.vmap(take_one_sample)(jax.random.split(rng_key, num_samples))

samples = prior_sample(rng_key, 10, Y_rv)
samples

prior_sampler(Y_rv, mu_rv).run(rng_key, 1000, {a_tt: 1.})