<jupyter_start><jupyter_text>Use BlackJAX with PyMC v4Author: <NAME>, based on the notebook for PyMC3 by <NAME><jupyter_code>import jax
import pymc as pm
from pymc.sampling_jax import get_jaxified_logp
import numpy as np
import blackjax.nuts as nuts
import blackjax.stan_warmup as stan_warmup

# Just for me to disable the GPU
import os
os.environ['CUDA_VISIBLE_DEVICES'] = ''

print(f"Running on PyMC3 v{pm.__version__}")<jupyter_output>You are running the v4 development version of PyMC which currently still lacks key features. You probably want to use the stable v3 instead which you can either install via conda or find on the v3 GitHub branch: https://github.com/pymc-devs/pymc/tree/v3<jupyter_text>DataPlease refer to the [original TFP example](https://www.tensorflow.org/probability/examples/Eight_Schools) for a description of the problem and the model that is used.<jupyter_code># Data of the Eight Schools Model
J = 8
y = np.array([28.0, 8.0, -3.0, 7.0, -1.0, 1.0, 18.0, 12.0])
sigma = np.array([15.0, 10.0, 16.0, 11.0, 9.0, 11.0, 10.0, 18.0])<jupyter_output><empty_output><jupyter_text>Model<jupyter_code>with pm.Model() as model:

    mu = pm.Normal("mu", mu=0.0, sigma=10.0)
    tau = pm.HalfCauchy("tau", 5.0)

    theta = pm.Normal("theta", mu=0, sigma=1, shape=J)
    theta_1 = mu + tau * theta
    obs = pm.Normal("obs", mu=theta, sigma=sigma, shape=J, observed=y)
chains = 4
n_warmup = 1000
n_samples = 1_000<jupyter_output><empty_output><jupyter_text>Sampling using PyMC NUTS Sampler<jupyter_code>%%time

with model:
    posterior = pm.sample(draws=n_samples, tune=n_warmup, chains=chains)<jupyter_output>Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
/home/martin/miniconda3/envs/dl/lib/python3.8/site-packages/pymc/model.py:925: FutureWarning: `Model.initial_point` has been deprecated. Use `Model.recompute_initial_point(seed=None)`.
  warnings.warn(
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [mu, tau, theta]
/home/martin/miniconda3/envs/dl/lib/python3.8/site-packages/pymc/model.py:925: FutureWarning: `Model.initial_point` has been deprecated. Use `Model.recompute_initial_point(seed=None)`.
  warnings.warn(<jupyter_text>Sampling using PyMC JAX Numpyro NUTS sampler<jupyter_code>%%time

with model:
    hierarchical_trace_jax = pm.sampling_jax.sample_numpyro_nuts(
        n_samples, target_accept=0.9, chains=chains, tune=n_warmup 
    )<jupyter_output>Compiling...<jupyter_text>Sampling using BlackJax Configuring the model for BlackJax<jupyter_code>seed = jax.random.PRNGKey(1234)

logp_fn_jax = get_jaxified_logp(model)

rv_names = [rv.name for rv in model.value_vars]
initial_point = model.recompute_initial_point()
init_state = [initial_point[rv_name] for rv_name in rv_names]
init_state_batched = jax.tree_map(lambda x: np.repeat(x[None, ...], chains, axis=0), init_state)

initial_state = jax.vmap(nuts.new_state, in_axes=(0, None))(init_state_batched, logp_fn_jax)
%%time

# Warmup
kernel_factory = lambda step_size, inverse_mass_matrix: nuts.kernel(
    logp_fn_jax, step_size, inverse_mass_matrix
)

init_seeds = jax.random.split(seed, 4)

warmup_fn = lambda seed, init_state: stan_warmup.run(seed, kernel_factory, init_state, n_warmup)

last_state, (step_size, inverse_mass_matrix), _ = jax.vmap(warmup_fn)(init_seeds, initial_state)
%%time

# Sampling
def inner_inference_loop(
    rng_key, step_size, inverse_mass_matrix, initial_state, num_samples, kernel_factory
):
    
    kernel = kernel_factory(step_size, inverse_mass_matrix)
    
    def one_step(state, rng_key):
        
        key = jax.random.split(rng_key, 2)
        state, info = kernel(key, state)
        return state, (state, info)
    
    keys = jax.random.split(rng_key, num_samples)
    _, states = jax.lax.scan(one_step, initial_state, keys)
    
    return states

from functools import partial

to_vmap = partial(inner_inference_loop, num_samples=n_samples, kernel_factory=kernel_factory)

# TODO: Should I be doing this?
seed = jax.random.PRNGKey(12345)
start_keys = jax.random.split(seed, 4)

results = jax.vmap(to_vmap)(start_keys, step_size, inverse_mass_matrix, last_state)
step_size
inverse_mass_matrix[3]
draw_list = results[0].position

draw_dict = {x: y for x, y in zip(rv_names, draw_list)}
import arviz as az
draws_untransformed_arviz = az.from_dict(draw_dict)

# Result I'm getting for theta
az.plot_posterior(draws_untransformed_arviz.posterior['theta'])
# Result we probably should be getting
az.plot_posterior(posterior.posterior['theta'])
# Additional work to transform variables (not important here, but seems to work)
from aesara.graph.fg import FunctionGraph
from pymc.util import get_default_varnames
from aesara.graph import optimize_graph
from aesara.link.jax.dispatch import jax_funcify

def transform_raw_samples(model, draw_list, keep_untransformed=False):
    
    var_names = model.unobserved_value_vars
    vars_to_sample = list(get_default_varnames(var_names, include_transformed=keep_untransformed))
    
    mcmc_samples = {}
    
    for v in vars_to_sample:
        fgraph = FunctionGraph(model.value_vars, [v], clone=False)
        optimize_graph(fgraph, include=["fast_run"], exclude=["cxx_only", "BlasOpt"])
        jax_fn = jax_funcify(fgraph)
        result = jax.vmap(jax.vmap(jax_fn))(*draw_list)[0]
        mcmc_samples[v.name] = result
        
    return mcmc_samples
mcmc_samples = transform_raw_samples(model, draw_list)
draws_arviz = az.from_dict(mcmc_samples)
# Result I'm getting for theta
az.plot_posterior(draws_arviz.posterior['theta'])
# Result we probably should be getting
az.plot_posterior(posterior.posterior['theta'])<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10712391727063815, 'fraction_numerical': 0.018030758352483647, 'mean_word_length': 5.3006681514476615, 'pattern_counts': {'":': 0, '<': 24, '<?xml version=': 0, '>': 24, 'https://': 2, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '447042', 'license': 'Apache-2.0', 'n_tokens_mistral': 2155, 'n_tokens_neox': 2045, 'n_words': 452, 'path': 'notebooks/Use with PyMC v4.ipynb', 'repo_name': 'martiningram/blackjax', 'size': '5676'}