### Set up Eight Schools Example Data

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_use_pymc.md

Initializes the dataset for the Eight Schools example, including observed values and standard deviations.

```python
import numpy as np


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])
```

--------------------------------

### Install uv

Source: https://github.com/blackjax-devs/blackjax/blob/main/CONTRIBUTING.md

Use this command to install the uv package manager.

```bash
curl -LsSf https://astral.sh/uv/install.sh | sh
```

--------------------------------

### NUTS Sampler Example in BlackJAX

Source: https://github.com/blackjax-devs/blackjax/blob/main/README.md

A self-contained example demonstrating how to use the NUTS (No-U-Turn Sampler) from BlackJAX. It requires JAX and NumPy. Ensure JAX is installed with GPU support if needed.

```python
import jax
import jax.numpy as jnp
import jax.scipy.stats as stats
import numpy as np

import blackjax

observed = np.random.normal(10, 20, size=1_000)
def logdensity_fn(x):
    logpdf = stats.norm.logpdf(observed, x["loc"], x["scale"])
    return jnp.sum(logpdf)

# Build the kernel
step_size = 1e-3
inverse_mass_matrix = jnp.array([1., 1.])
nuts = blackjax.nuts(logdensity_fn, step_size, inverse_mass_matrix)

# Initialize the state
initial_position = {"loc": 1., "scale": 2.}
state = nuts.init(initial_position)

# Iterate
rng_key = jax.random.key(0)
step = jax.jit(nuts.step)
for i in range(100):
    nuts_key = jax.random.fold_in(rng_key, i)
    state, _ = step(nuts_key, state)
```

--------------------------------

### Setup and Imports

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_use_funsor.md

Import necessary libraries and configure the Funsor backend to JAX.

```python
import jax
import jax.numpy as jnp
import numpy as np
import blackjax

import funsor
import funsor.ops as ops

funsor.set_backend("jax")          # must be called before Tensor/Variable are used

from funsor.domains import Bint
from funsor.tensor import Tensor
from funsor.terms import Variable
from funsor.jax.distributions import Categorical, Normal

from datetime import date
rng_key = jax.random.key(int(date.today().strftime("%Y%m%d")))
```

--------------------------------

### Setup Eight Schools Example Data

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_use_numpyro.md

Initializes the JAX random number generator and defines the dataset for the Eight Schools example.

```python
import numpy as np


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])
```

```python
import jax

from datetime import date
rng_key = jax.random.key(int(date.today().strftime("%Y%m%d")))
```

--------------------------------

### Install PyTensor

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_use_pytensor.md

Install PyTensor using pip. This is a prerequisite for running the examples in this document.

```bash
pip install pytensor
```

--------------------------------

### Convert example to Jupyter notebook

Source: https://github.com/blackjax-devs/blackjax/blob/main/CONTRIBUTING.md

Converts a Markdown example file into a Jupyter notebook format using Jupytext.

```shell
jupytext docs/examples/your_example_file.md --to notebook
```

--------------------------------

### Sync development environment

Source: https://github.com/blackjax-devs/blackjax/blob/main/CONTRIBUTING.md

Creates a local virtual environment and installs all development dependencies.

```bash
uv sync --group dev --extra progress
```

--------------------------------

### Install Funsor

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_use_funsor.md

Install the required Funsor library version.

```bash
pip install "funsor>=0.4.7"
```

--------------------------------

### Import Libraries and Setup

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_use_tfp.md

Imports necessary libraries including JAX, NumPy, and TensorFlow Probability. Sets up a random number generator key.

```python
import numpy as np


num_schools = 8  # number of schools
treatment_effects = np.array(
    [28, 8, -3, 7, -1, 1, 18, 12], dtype=np.float32
)  # treatment effects
treatment_stddevs = np.array(
    [15, 10, 16, 11, 9, 11, 10, 18], dtype=np.float32
)  # treatment SE
```

```python
import jax
import jax.numpy as jnp

from datetime import date
rng_key = jax.random.key(int(date.today().strftime("%Y%m%d")))
```

--------------------------------

### Setup BlackJAX and Dependencies

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/diagnostics.md

Imports necessary libraries for MCMC sampling and diagnostics.

```python
import jax
import jax.numpy as jnp
import jax.scipy.stats as stats
import matplotlib.pyplot as plt
import numpy as np
import arviz as az

import blackjax
```

--------------------------------

### Install BlackJAX

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/index.md

Commands to install the BlackJAX library using standard package managers.

```bash
pip install blackjax
```

```bash
conda install blackjax -c conda-forge
```

--------------------------------

### Install BlackJAX with pip

Source: https://github.com/blackjax-devs/blackjax/blob/main/README.md

Install the BlackJAX library using pip. This is the standard method for most Python environments.

```bash
pip install blackjax
```

--------------------------------

### Install Dependencies

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_use_funsor.md

Required packages for using NumPyro with Funsor.

```bash
pip install numpyro "funsor>=0.4.7"
```

--------------------------------

### Initializing Multiple Chains

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_sample_multiple_chains.md

Prepare initial states for multiple chains by vectorizing the `init` function using `jax.vmap`. This ensures each chain starts with its own unique initial configuration.

```python
initial_positions = {"loc": np.ones(num_chains), "log_scale": np.ones(num_chains)}
initial_states = jax.vmap(nuts.init, in_axes=(0))(initial_positions)
```

--------------------------------

### Numpyro Model Example

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_use_numpyro.md

This snippet shows a basic Numpyro model definition. Ensure all necessary imports are present before running.

```python
import numpyro
import numpyro.distributions as dist
import jax.numpy as jnp

def model(data=None):
    # Hyperpriors
    mu = numpyro.sample("mu", dist.Normal(0, 1))
    sigma = numpyro.sample("sigma", dist.HalfCauchy(1))

    # Likelihood
    with numpyro.plate("data_loop", data.shape[0]):
        numpyro.sample("obs", dist.Normal(mu, sigma), obs=data)

```

--------------------------------

### Setup RNG Key

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_custom_gradients.md

Initializes the random number generator key for JAX operations.

```python
import jax

from datetime import date
rng_key = jax.random.key(int(date.today().strftime("%Y%m%d")))
```

--------------------------------

### Setup Plotting and Load Image

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_reproduce_the_blackjax_image.md

Configures plotting parameters and loads the image into a NumPy array for processing. Ensures the image is converted to a mask representing black pixels.

```python
import matplotlib
import matplotlib.image as mpimg
import matplotlib.pyplot as plt
import numpy as np

plt.rcParams["axes.spines.right"] = False
plt.rcParams["axes.spines.top"] = False
matplotlib.rcParams['animation.embed_limit'] = 25
```

```python
# Load the image
im = mpimg.imread('./data/blackjax.png')

# Convert to mask
im = np.amax(im[:, :, :2], 2) < 0.5

# Convert back to float
im = im.astype(float)
```

--------------------------------

### Sample Initial Weights for BNN

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_use_oryx.md

Builds a Bayesian Neural Network (BNN) and samples initial weights using `joint_sample`. This provides a starting point for inference algorithms.

```python
from oryx.core.ppl import joint_sample

# Assuming bnn, num_features, and rng_key are defined elsewhere
# bnn = mlp([50, 50], num_classes)
# rng_key, init_key = jax.random.split(rng_key)
# initial_weights = joint_sample(bnn)(init_key, jnp.ones(num_features))

# print(initial_weights.keys())
```

--------------------------------

### HMC Kernel initialization and step with Optax design

Source: https://github.com/blackjax-devs/blackjax/wiki/Meeting-minutes.org

Demonstrates the Optax-inspired design for HMC kernel initialization and stepping, using a class instance with init and step methods.

```python
import blackjax

class SamplingKernel(NamedTuple):
    init: Callable
    step: Callable

# Now
init, step = bjx.hmc(logprob_fn)
state = init(position)
new_state, info = step(rng_key, state)

# With Optax design
hmc = blackjax.hmc(logprob_fn)
state = hmc.init(position)
new_state, info = hmc.step(rng_key, state)
```

--------------------------------

### Install BlackJAX with conda-forge

Source: https://github.com/blackjax-devs/blackjax/blob/main/README.md

Install the BlackJAX library using conda-forge. This is an alternative installation method for users who prefer conda environments.

```bash
conda install -c conda-forge blackjax
```

--------------------------------

### Build documentation

Source: https://github.com/blackjax-devs/blackjax/blob/main/CONTRIBUTING.md

Generates the HTML documentation using Sphinx.

```bash
make build-docs
```

--------------------------------

### Import required libraries

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/quickstart.md

Initialize the environment by importing JAX, NumPy, and BlackJAX modules.

```python
import matplotlib.pyplot as plt
import numpy as np

import jax
import jax.numpy as jnp
import jax.scipy.stats as stats

import blackjax
```

--------------------------------

### Install ipywidgets for Jupyter

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_progress_bar.md

Install the ipywidgets package to enable rich progress bar widgets in Jupyter notebook environments.

```bash
pip install ipywidgets
```

--------------------------------

### Run NUTS Window Adaptation

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_use_funsor.md

Tunes the NUTS step size and mass matrix over 1000 warmup steps.

```python
%%time

rng_key, warmup_key = jax.random.split(rng_key)

adapt = blackjax.window_adaptation(blackjax.nuts, gmm_logdensity)
(last_state, parameters), _ = adapt.run(warmup_key, position0, num_steps=1000)
kernel = blackjax.nuts(gmm_logdensity, **parameters).step
```

--------------------------------

### Numpydoc Style Docstring Example

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/developer/design_principles.md

Example of a function docstring following the numpydoc format, including Parameters and Returns sections. Magic numbers should be explained.

```python
def build_kernel(step_size: float, inverse_mass_matrix: Array) -> Callable:
    """Build an HMC transition kernel.

    Parameters
    ----------
    step_size
        Size of the leapfrog integration step.
    inverse_mass_matrix
        Inverse of the mass matrix. Either a 1D array (diagonal) or 2D array (dense).

    Returns
    -------
    A kernel function ``kernel(rng_key, state, logdensity_fn) -> (HMCState, HMCInfo)``.
    """
```

--------------------------------

### Run a NUTS sampler in Blackjax

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/index.md

Demonstrates defining a log-density function, initializing a NUTS kernel, and running the sampling loop.

```Python
import jax
import jax.numpy as jnp
import jax.scipy.stats as stats
import numpy as np

import blackjax

observed = np.random.normal(10, 20, size=1_000)
def logdensity_fn(x):
    logpdf = stats.norm.logpdf(observed, x["loc"], x["scale"])
    return jnp.sum(logpdf)

# Build the kernel
step_size = 1e-3
inverse_mass_matrix = jnp.array([1., 1.])
nuts = blackjax.nuts(logdensity_fn, step_size, inverse_mass_matrix)

# Initialize the state
initial_position = {"loc": 1., "scale": 2.}
state = nuts.init(initial_position)

# Iterate
rng_key = jax.random.key(0)
step = jax.jit(nuts.step)
for i in range(1_000):
    nuts_key = jax.random.fold_in(rng_key, i)
    state, _ = step(nuts_key, state)
```

--------------------------------

### Import necessary libraries

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_other_frameworks.md

Import JAX, JAX NumPy, NumPy, and date for setting up the random number generator.

```python
import jax
import jax.numpy as jnp
import numpy as np
from datetime import date
rng_key = jax.random.key(int(date.today().strftime("%Y%m%d")))
```

--------------------------------

### Run Window Adaptation for NUTS Parameters

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_use_oryx.md

Uses `blackjax.window_adaptation` to find optimal parameters for the NUTS algorithm by running a warmup phase. This improves sampling efficiency.

```python
%%time
import blackjax

# Assuming rng_key, logdensity_fn, and initial_weights are defined
# rng_key, warmup_key = jax.random.split(rng_key)
# adapt = blackjax.window_adaptation(blackjax.nuts, logdensity_fn)
# (last_state, parameters), _ = adapt.run(warmup_key, initial_weights, 100)
# kernel = blackjax.nuts(logdensity_fn, **parameters).step
```

--------------------------------

### Initialize HMC State

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/quickstart.md

Use the hmc.init function to generate the initial state from a starting position dictionary.

```python
initial_position = {"loc": 1.0, "log_scale": 1.0}
initial_state = hmc.init(initial_position)
initial_state
```

--------------------------------

### Configure and Initialize Tempered SMC Sampler

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_reproduce_the_blackjax_image.md

Sets up the tempered SMC sampler using a Random Walk Metropolis-Hastings (RWMH) kernel. This involves defining the temperature schedule, the proposal distribution, and initializing the SMC state with initial samples.

```python
import functools
import blackjax
import blackjax.smc.resampling as resampling

# Temperature schedule
n_temperatures = 150
lambda_schedule = np.logspace(-3, 0, n_temperatures)

# The proposal distribution is a random walk with a fixed scale
scale = 0.5
normal = blackjax.mcmc.random_walk.normal(scale * jnp.ones((2,)))

# Bake the proposal into the kernel so no function leaks into mcmc_parameters
rw_kernel = functools.partial(
    blackjax.additive_step_random_walk.build_kernel(),
    random_step=normal,
)
rw_init = blackjax.additive_step_random_walk.init

tempered = blackjax.tempered_smc(
    prior_logpdf,
    log_likelihood,
    rw_kernel,
    rw_init,
    {},   # mcmc_parameters must be JAX-traceable; proposal is baked into rw_kernel
    resampling.systematic,
    num_mcmc_steps=5,
)

initial_smc_state = tempered.init(zs_init)
```

--------------------------------

### Get JAXified Log-Probability Function

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_use_pymc.md

Extracts a JAX-compatible log-probability function from the PyMC model using `get_jaxified_logp`.

```python
from pymc.sampling.jax import get_jaxified_logp

rvs = [rv.name for rv in model.value_vars]
logdensity_fn = get_jaxified_logp(model)
```

--------------------------------

### Run pre-commit and tests

Source: https://github.com/blackjax-devs/blackjax/blob/main/CONTRIBUTING.md

Execute these commands to ensure code quality and pass tests before pushing changes.

```bash
uv run pre-commit run --all-files
make test
```

--------------------------------

### HMCParameters NamedTuple for kernel parameters

Source: https://github.com/blackjax-devs/blackjax/wiki/Meeting-minutes.org

Example of using a NamedTuple to structure parameters for the HMC kernel, including optional fields.

```python
class HMCParameters(NamedTuple):
    num_integration_steps: Optional[int]
    step_size: Optional[float]
    inverse_mass_matrix: Optional[jnp.ndarray]

params = HMCParameters(10)
```

--------------------------------

### Create Top-Level VI API Wrapper

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/developer/new_algorithm_guide.md

Constructs a user-facing API for a Variational Inference algorithm using `VIAlgorithm`. This wrapper defines `init_fn`, `step_fn`, and `sample_fn` for the VI process.

```python
from blackjax.base import VIAlgorithm

def as_top_level_api(
    logdensity_fn: Callable,
    optimizer,       # optax GradientTransformation
    num_samples: int = 100,
) -> VIAlgorithm:
    def init_fn(position):
        return init(position, logdensity_fn)

    def step_fn(rng_key, state):
        return step(rng_key, state, logdensity_fn, optimizer, num_samples)

    def sample_fn(rng_key, state, num_samples):
        return sample(rng_key, state, num_samples)

    return VIAlgorithm(init_fn, step_fn, sample_fn)
```

--------------------------------

### Define Bayesian linear regression model

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_progress_bar.md

Sets up the data, priors, and log-density function for a Bayesian linear regression problem used in sampling examples.

```python
import warnings

import jax
import jax.numpy as jnp
from jax.scipy.stats import norm

import blackjax
from blackjax.util import run_inference_algorithm

# Headless doc builds have no ipywidgets; the plain-text bar is intended here.
warnings.filterwarnings("ignore", message="IProgress not found")

N, DIM, SIGMA = 300, 10, 0.5
kx, kw, kn = jax.random.split(jax.random.key(0), 3)
X = jax.random.normal(kx, (N, DIM))
true_w = 2.0 * jax.random.normal(kw, (DIM,))
y = X @ true_w + SIGMA * jax.random.normal(kn, (N,))


def logprior(w):
    return norm.logpdf(w, 0.0, 5.0).sum()


def loglikelihood(w):
    return norm.logpdf(y - X @ w, 0.0, SIGMA).sum()


logdensity = lambda w: logprior(w) + loglikelihood(w)
```

--------------------------------

### Create JAX Mesh and NamedSharding

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_sample_multiple_chains.md

Set up a JAX `Mesh` to name device axes and define `NamedSharding` for distributing computations. This is the first step in preparing for device-parallel execution.

```python
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P

mesh = jax.make_mesh((num_chains,), ('chain',))
sharding = NamedSharding(mesh, P('chain'))
```

--------------------------------

### Define Model and Marginal for Laplace HMC

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_laplace_hmc.md

Setup the latent dimension, generate synthetic data, and define the log-joint and exact marginal functions for verification.

```python
n = 20  # latent dimension

# Generate data
true_log_sigma = jnp.log(jnp.array(2.0))
rng_key, data_key = jax.random.split(rng_key)
true_theta = jax.random.normal(data_key, (n,)) * jnp.exp(true_log_sigma)
rng_key, obs_key = jax.random.split(rng_key)
y_obs = true_theta + jax.random.normal(obs_key, (n,))


def log_joint(theta, log_sigma):
    """log p(theta, log_sigma, y).  theta is the latent, log_sigma is phi."""
    sigma = jnp.exp(log_sigma)
    log_prior_phi = stats.norm.logpdf(log_sigma, 0.0, 3.0)
    log_prior_theta = stats.norm.logpdf(theta, 0.0, sigma).sum()
    log_lik = stats.norm.logpdf(y_obs, theta, 1.0).sum()
    return log_prior_phi + log_prior_theta + log_lik


def exact_log_marginal(log_sigma):
    """Closed-form marginal p̂(log_sigma | y) for verification."""
    sigma = jnp.exp(log_sigma)
    log_prior_phi = stats.norm.logpdf(log_sigma, 0.0, 3.0)
    var_marg = sigma**2 + 1.0
    log_lik_marg = stats.norm.logpdf(y_obs, 0.0, jnp.sqrt(var_marg)).sum()
    return log_prior_phi + log_lik_marg
```

--------------------------------

### Initialize JAX and RNG Key

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_use_pymc.md

Sets up JAX for numerical operations and initializes a random number generator key for reproducibility.

```python
import jax
import jax.numpy as jnp

from datetime import date
rng_key = jax.random.key(int(date.today().strftime("%Y%m%d")))
```

--------------------------------

### Executing the General Sampling Loop

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_metropolis_within_gibbs.md

This code snippet demonstrates how to execute the general sampling loop with specific parameters, including the PRNG key, initial state, log-density function, stepping functions, initializers, and the number of samples. It also includes a timing decorator.

```python
%%time
positions_general = sampling_loop_general(
    rng_key=sample_key,  # reuse PRNG key from above
    initial_state=initial_state,
    logdensity_fn=logdensity,
    step_fn={
        "x": mwg_step_fn_x,
        "y": mwg_step_fn_y
    },
    init={
        "x": mwg_init_x,
        "y": mwg_init_y
    },
    parameters=parameters,
    num_samples=10_000
)
```

--------------------------------

### Run Stan Window Adaptation

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/quickstart.md

Automatically determine step size and inverse mass matrix using window adaptation.

```python
%%time

warmup = blackjax.window_adaptation(blackjax.nuts, logdensity)
rng_key, warmup_key, sample_key = jax.random.split(rng_key, 3)
(state, parameters), _ = warmup.run(warmup_key, initial_position, num_steps=1000)
```

--------------------------------

### Call and Differentiate Numba LogPDF in JAX

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_other_frameworks.md

Demonstrates calling the JAX-wrapped Numba log-probability density function (`numba_logpdf`) within JIT-compiled code and computing its gradient. This verifies the custom VJP setup.

```python
print(jax.jit(numba_logpdf)(1.0))
print(jax.grad(numba_logpdf)(1.0))
```

--------------------------------

### Sample from chains with adapted parameters

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_progress_bar.md

Performs parallel sampling using adapted parameters obtained from the warmup phase. The progress bar tracks the sampling process across all chains.

```python
k_sm = jax.random.split(jax.random.key(12), n_chains)


def sample_one(key, st, p):
    alg = blackjax.nuts(logdensity, **p)  # per-chain step_size + mass matrix
    _, hist = run_inference_algorithm(
        rng_key=key,
        initial_state=st,
        inference_algorithm=alg,
        num_steps=1500,
        transform=lambda s, _: s.position,
    )
    return hist


with blackjax.progress_bar(label=f"vmap sampling ({n_chains} chains)"):
    mc_chains = jax.vmap(sample_one)(k_sm, warm_states, chain_params)

print("Output shape (chains, steps, dim):", mc_chains.shape)
```

--------------------------------

### Run Laplace HMC Inference

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_laplace_hmc.md

Initialize the sampler with the log-joint and latent variable hints, then execute the inference algorithm.

```python
theta_init = jnp.zeros(n)
phi_init = jnp.array(0.0)

sampler = blackjax.laplace_hmc(
    log_joint,
    theta_init=theta_init,
    step_size=0.3,
    inverse_mass_matrix=jnp.ones(1),
    num_integration_steps=10,
    maxiter=100,  # L-BFGS iterations per leapfrog step
)

rng_key, run_key = jax.random.split(rng_key)
final_state, (states, infos) = run_inference_algorithm(
    run_key,
    sampler,
    num_steps=1_000,
    initial_position=phi_init,
    transform=lambda state, info: (state, info),
)

phi_samples = states.position          # shape (1000,)
theta_star_samples = states.theta_star # shape (1000, n) — MAP theta at each phi
```

--------------------------------

### Shard RNG Keys and Initial States

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_sample_multiple_chains.md

Distribute random number generator keys and initial states across devices using `jax.device_put` with the defined sharding. This ensures each chain starts with its own unique set of parameters.

```python
rng_key, sample_key = jax.random.split(rng_key)
sample_keys = jax.device_put(jax.random.split(sample_key, num_chains), sharding)
initial_states_sharded = jax.device_put(initial_states, sharding)
```

--------------------------------

### Initialize JAX Random Key and Suppress Warnings

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_use_oryx.md

Initializes a JAX random number generator key based on the current date and suppresses deprecation warnings from Oryx internals. This setup is necessary before proceeding with JAX and Oryx operations.

```python
import warnings
import jax
import jax.numpy as jnp

from datetime import date
rng_key = jax.random.key(int(date.today().strftime("%Y%m%d")))

# Suppress DeprecationWarnings from oryx internals (FlatPrimitive DebugInfo)
```

--------------------------------

### Initialize NUTS Sampler

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_sample_multiple_chains.md

Initializes the No-U-Turn Sampler (NUTS) with specified inverse mass matrix and step size. Requires the log-density function and Blackjax library.

```python
import blackjax


inv_mass_matrix = np.array([0.5, 0.01])
step_size = 1e-3

nuts = blackjax.nuts(logdensity, step_size, inv_mass_matrix)
```

--------------------------------

### Execute Inference and Process Samples

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/quickstart.md

Run the inference loop and transform the resulting samples into the desired parameter space.

```python
%%time
rng_key, sample_key = jax.random.split(rng_key)
states = inference_loop(sample_key, hmc_kernel, initial_state, 10_000)

mcmc_samples = states.position
mcmc_samples["scale"] = jnp.exp(mcmc_samples["log_scale"]).block_until_ready()
```

--------------------------------

### Run benchmark test

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/speed_up_guide.md

Execute the specific benchmark test for comparing flat versus dictionary parameterization.

```bash
JAX_PLATFORM_NAME=cpu pytest tests/test_benchmarks.py::test_horseshoe_nuts_flat_vs_dict -v -s
```

--------------------------------

### Initialize Laplace HMC environment

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_laplace_hmc.md

Imports necessary libraries and initializes the random number generator key for the Laplace marginalization workflow.

```python
from datetime import date

import jax
import jax.numpy as jnp
import jax.scipy.stats as stats
import matplotlib.pyplot as plt
import numpy as np

import blackjax
from blackjax.mcmc.laplace_marginal import laplace_marginal_factory
from blackjax.util import run_inference_algorithm

rng_key = jax.random.key(int(date.today().strftime("%Y%m%d")))
```

--------------------------------

### Algorithm State Initialization

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/developer/design_principles.md

The init function creates the initial algorithm state. It accepts position, logdensity_fn, and an optional rng_key.

```python
init(position, logdensity_fn, *, rng_key=None) -> State
```

--------------------------------

### Import NumPyro and Funsor Utilities

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_use_funsor.md

Necessary imports for defining and initializing NumPyro models with Funsor support.

```python
import numpyro
import numpyro.distributions as dist
from numpyro.contrib.funsor import config_enumerate
from numpyro.infer.util import initialize_model
```

--------------------------------

### Display a labeled progress bar for sampling

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_progress_bar.md

Wrap a sampling call with the progress_bar context manager to enable live updates during execution.

```python
warmup = blackjax.window_adaptation(blackjax.nuts, logdensity)
k_warm = jax.random.key(1)

with blackjax.progress_bar(label="NUTS warmup"):
    (state, params), _ = warmup.run(k_warm, jnp.zeros(DIM), num_steps=1000)
```

--------------------------------

### Run Window Adaptation

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_use_funsor.md

Perform window adaptation for NUTS using the defined log-density function.

```python
%%time

rng_key, warmup_key_npy = jax.random.split(rng_key)

adapt_npy = blackjax.window_adaptation(blackjax.nuts, logdensity_fn_npy)
(last_state_npy, params_npy), _ = adapt_npy.run(
    warmup_key_npy, initial_position_npy, num_steps=1000
)
kernel_npy = blackjax.nuts(logdensity_fn_npy, **params_npy).step
```

--------------------------------

### Define Initial Parameter Function for Blackjax

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_use_pytensor.md

Define a function to generate initial positions for the sampler. This function takes a JAX random key and returns a dictionary of parameters.

```python
import blackjax

def init_param_fn(seed):
    key1, key2, key3 = jax.random.split(seed, 3)
    return {
        "log_a":       jax.random.normal(key1, ())
        "log_b":       jax.random.normal(key2, ())
        "logit_theta": jax.random.normal(key3, (n_rat_tumors,))
    }
```

--------------------------------

### Configure HMC Sampler Parameters

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/quickstart.md

Define the inverse mass matrix, integration steps, and step size before initializing the HMC sampler.

```python
inv_mass_matrix = np.array([0.5, 0.01])
num_integration_steps = 60
step_size = 1e-3

hmc = blackjax.hmc(logdensity, step_size, inv_mass_matrix, num_integration_steps)
```

--------------------------------

### Adapt NUTS Kernel with Window Adaptation

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_use_numpyro.md

Performs window adaptation for the NUTS sampler using Blackjax to tune sampler parameters based on the log-probability function and initial position.

```python
import blackjax

num_warmup = 2000

adapt = blackjax.window_adaptation(
    blackjax.nuts, logdensity_fn, target_acceptance_rate=0.8
)
rng_key, warmup_key = jax.random.split(rng_key)
(last_state, parameters), _ = adapt.run(warmup_key, initial_position, num_warmup)
```

--------------------------------

### Generic Sampling Algorithm Interface

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/developer/new_algorithm_guide.md

Illustrates the typical user-facing interface for BlackJAX sampling algorithms like MCMC. Requires an initializer and an iterator.

```python
sampling_algorithm = blackjax.nuts(logdensity_fn, step_size, inverse_mass_matrix)
state = sampling_algorithm.init(initial_position)
new_state, info = sampling_algorithm.step(rng_key, state)
```

--------------------------------

### Running the Sampler

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_sample_multiple_chains.md

Execute the MCMC sampler for a specified number of samples and chains. This snippet includes timing and blocking until computation is ready.

```python
rng_key, sample_key = jax.random.split(rng_key)
states = inference_loop_multiple_chains(
    sample_key, nuts.step, initial_states, 2_000, num_chains
)
_ = states.position["loc"].block_until_ready()
```

--------------------------------

### Verifying JAX Device Count

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_sample_multiple_chains.md

Confirm that JAX has successfully recognized multiple CPU devices after configuring the `XLA_FLAGS` environmental variable.

```python
len(jax.devices())
```

--------------------------------

### Adapt NUTS Sampler with BlackJAX

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_use_pymc.md

Performs window adaptation for the NUTS sampler using BlackJAX, tuning parameters based on the log-probability function and initial position.

```python
import blackjax
# Get the initial position from PyMC
init_position_dict = model.initial_point()
init_position = [init_position_dict[rv] for rv in rvs]

rng_key, warmup_key = jax.random.split(rng_key)

adapt = blackjax.window_adaptation(blackjax.nuts, logdensity_fn)
(last_state, parameters), _ = adapt.run(warmup_key, init_position, 1000)
kernel = blackjax.nuts(logdensity_fn, **parameters).step
```

--------------------------------

### Run inference using the standard algorithm wrapper

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/speed_up_guide.md

Use run_inference_algorithm to automatically handle jax.jit and jax.lax.scan for efficient sampling.

```python
import blackjax
from blackjax.util import run_inference_algorithm

nuts = blackjax.nuts(logdensity_fn, step_size, inv_mass_matrix)

final_state, (states, infos) = run_inference_algorithm(
    rng_key,
    nuts,
    num_steps=1_000,
    initial_position=initial_position,
    transform=lambda state, info: (state.position, info),
)
```

--------------------------------

### Handle JIT cache with progress bars

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_progress_bar.md

Demonstrates how JIT-compiled functions interact with progress bar contexts when compiled before or during context entry.

```python
sample_jitted = jax.jit(
    lambda key: run_inference_algorithm(
        rng_key=key,
        initial_state=state,
        inference_algorithm=nuts,
        num_steps=100,
        transform=lambda st, _: st.position,
    )[1]
)

# Context 1: traces the function, bakes _dynamic_router
with blackjax.progress_bar(label="chain 1 (traces)"):
    _ = sample_jitted(jax.random.key(7))

# Context 2: cache hit → _dynamic_router routes to ctx2's recorder → bar works
with blackjax.progress_bar(label="chain 2 (cache hit → bar works)"):
    _ = sample_jitted(jax.random.key(8))
```

```python
import warnings

# Compile entirely outside any progress_bar context.
jax.clear_caches()
pre_compiled = jax.jit(
    lambda key: run_inference_algorithm(
        rng_key=key,
        initial_state=state,
        inference_algorithm=nuts,
        num_steps=100,
        transform=lambda st, _: st.position,
    )[1]
)
_ = pre_compiled(jax.random.key(10))  # compiled now, no callback baked in

# Inside a context: 0 events → warning fires.
with warnings.catch_warnings(record=True) as w:
    warnings.simplefilter("always")
    with blackjax.progress_bar(label="pre-compiled (no bar)"):
        _ = pre_compiled(jax.random.key(11))
    if w:
        print("Warning (expected):", w[0].message)

# Fix: clear caches to force retrace with the callback baked in.
jax.clear_caches()
with blackjax.progress_bar(label="after clear_caches (bar works)"):
    _ = pre_compiled(jax.random.key(12))  # retraced, bar now works
```

--------------------------------

### Initialize NumPyro Model

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_use_funsor.md

Initialize the model to obtain the potential function and initial parameters.

```python
rng_key, init_key_npy = jax.random.split(rng_key)

init_params_npy, potential_fn_gen_npy, *_ = initialize_model(
    init_key_npy,
    gmm_model_npy,
    model_args=(data, K),
    dynamic_args=True,
)
initial_position_npy = init_params_npy.z
print("Sites and shapes:", {k: v.shape for k, v in initial_position_npy.items()})
```

--------------------------------

### Initialize Numpyro Model and Potential Function

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_use_numpyro.md

Uses Numpyro's `initialize_model` to obtain initial parameters and a potential function generator, which is then adapted for Blackjax.

```python
from numpyro.infer.util import initialize_model

rng_key, init_key = jax.random.split(rng_key)
init_params, potential_fn_gen, *_ = initialize_model(
    init_key,
    eight_schools_noncentered,
    model_args=(J, sigma, y),
    dynamic_args=True,
)

logdensity_fn = lambda position: -potential_fn_gen(J, sigma, y)(position)
initial_position = init_params.z
```

--------------------------------

### Create Top-Level API Wrapper

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/developer/new_algorithm_guide.md

Wraps a custom MCMC kernel into a user-facing API using `build_sampling_algorithm`. This simplifies initialization and stepping for users. Pass `pass_rng_key_to_init=True` if your initialization requires an RNG key.

```python
from blackjax.base import SamplingAlgorithm, build_sampling_algorithm

def as_top_level_api(
    logdensity_fn: Callable,
    step_size: float,
) -> SamplingAlgorithm:
    """My Sampler — user-facing convenience wrapper.

    Examples
    --------

    .. code::

        sampler = blackjax.my_sampler(logdensity_fn, step_size=0.1)
        state = sampler.init(initial_position)
        new_state, info = sampler.step(rng_key, state)

    Parameters
    ----------
    logdensity_fn
        The log-density function of the target distribution.
    step_size
        Proposal step size.

    Returns
    -------
    A ``SamplingAlgorithm``.
    """
    kernel = build_kernel()
    return build_sampling_algorithm(kernel, init, logdensity_fn,
                                    kernel_args=(step_size,))
```

--------------------------------

### Run Inference with Adapted Parameters

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/quickstart.md

Use the parameters obtained from adaptation to run the NUTS sampler.

```python
%%time

kernel = blackjax.nuts(logdensity, **parameters).step
states = inference_loop(sample_key, kernel, state, 1_000)

mcmc_samples = states.position
mcmc_samples["scale"] = jnp.exp(mcmc_samples["log_scale"]).block_until_ready()
```

--------------------------------

### Initialize and Animate Plot

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_reproduce_the_blackjax_image.md

Sets up Matplotlib for animation, configures plot aesthetics, and creates an animation from the BlackJAX samples. Ensure 'ys_init', 'xs_init', 'n_samples', and 'n_temperatures' are defined before use.

```python
import matplotlib.animation as animation
from IPython.display import HTML

plt.style.use('dark_background')
animation.embed_limit = 25
fig, ax = plt.subplots()
fig.tight_layout()

ax.set_axis_off()
ax.set_xlim(0, 250)
ax.set_ylim(0, 80)

scat = ax.scatter(ys_init, 80 - xs_init, s=1000 * 1 / n_samples)


def animate(i):
    scat.set_offsets(np.c_[samples[i, :, 1], 80 - samples[i, :, 0]])
    scat.set_sizes(1000 * weights[i])
    return scat,


ani = animation.FuncAnimation(fig, animate, repeat=True,
                              frames=n_temperatures, blit=True, interval=100)
HTML(ani.to_jshtml())
```

--------------------------------

### Initialize BlackJAX Kernels for MWG

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_metropolis_within_gibbs.md

Sets up the initialization functions for the RMH and HMC kernels, which will be used as components within the Metropolis-Within-Gibbs sampler.

```python
mwg_init_x = blackjax.rmh.init
mwg_init_y = blackjax.hmc.init
```

--------------------------------

### Verify GMM Log-Density and Gradients

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_use_funsor.md

Verifies the log-density calculation and gradient computation using jax.grad.

```python
position0 = {
    "mu":     jnp.array([-3., 0., 3.]),
    "log_pi": jnp.zeros(K - 1),   # K-1 free logits; first logit fixed at 0
}
print("log p(data | init):", gmm_logdensity(position0))
print("grad w.r.t. mu    :", jax.grad(gmm_logdensity)(position0)["mu"])
```

--------------------------------

### Configure XLA device count

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_progress_bar.md

Sets the XLA host platform device count to match the number of available CPU cores.

```python
import os
import multiprocessing

os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count={}".format(
    multiprocessing.cpu_count()
)
```

--------------------------------

### Execute Chains in Parallel with jax.shard_map and jax.jit

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_sample_multiple_chains.md

Compile and execute the `run_one_chain` function in parallel across devices using `jax.jit` and `jax.shard_map`. This significantly speeds up the sampling process by running chains concurrently.

```python
%%time
sharded_states = jax.jit(jax.shard_map(
    run_one_chain,
    mesh=mesh,
    in_specs=(P('chain'), P('chain')),
    out_specs=P('chain'),
    check_vma=False,
))(sample_keys, initial_states_sharded)
_ = sharded_states.position["loc"].block_until_ready()
```

--------------------------------

### Adapt NUTS Sampler with Window Adaptation

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_use_pytensor.md

Adapt the NUTS sampler using Blackjax's window adaptation. This process tunes sampler parameters based on an initial run.

```python
rng_key, init_key, warmup_key, sample_key = jax.random.split(rng_key, 4)
init_position = init_param_fn(init_key)

adapt = blackjax.window_adaptation(blackjax.nuts, logdensity_fn)
(state, parameters), _ = adapt.run(warmup_key, init_position, 1000)
kernel = blackjax.nuts(logdensity_fn, **parameters).step
```

--------------------------------

### Build HMC Kernel

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/quickstart.md

JIT-compile the HMC step function for improved performance.

```python
hmc_kernel = jax.jit(hmc.step)
```

--------------------------------

### Visualize NUTS Samples

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/quickstart.md

Plot the trace of the sampled parameters.

```python
:tags: [hide-input]

fig, (ax, ax1) = plt.subplots(ncols=2, figsize=(15, 6))
ax.plot(mcmc_samples["loc"])
ax.set_xlabel("Samples")
ax.set_ylabel("loc")

ax1.plot(mcmc_samples["scale"])
ax1.set_xlabel("Samples")
ax1.set_ylabel("scale");
```

--------------------------------

### Run Window Adaptation for HMC Parameters

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_use_tfp.md

Performs window adaptation to tune the step size and inverse mass matrix for the HMC integrator. Uses 3 integration steps per adaptation iteration.

```python
import blackjax


initial_position = {
    "avg_effect": jnp.zeros([]),
    "avg_stddev": jnp.zeros([]),
    "school_effects_standard": jnp.ones([num_schools]),
}


rng_key, warmup_key = jax.random.split(rng_key)
adapt = blackjax.window_adaptation(
    blackjax.hmc, logdensity_fn, num_integration_steps=3
)

(last_state, parameters), _ = adapt.run(warmup_key, initial_position, 1000)
kernel = blackjax.hmc(logdensity_fn, **parameters).step
```

--------------------------------

### Register Sampler in __init__.py

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/developer/new_algorithm_guide.md

Imports a custom sampler module into the main `__init__.py` file to make it accessible at the top level of the BlackJAX library.

```python
# At the top of __init__.py, import your module:
from .mcmc import my_sampler as _my_sampler
```

--------------------------------

### Sample from Posterior using Blackjax NUTS with Numba LogPDF

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_other_frameworks.md

This snippet shows how to use Blackjax's NUTS sampler with the custom `numba_logpdf` function. It initializes the sampler, performs MCMC steps, and prints the final state, demonstrating the integration of Numba-compiled models into probabilistic programming workflows.

```python
import blackjax

inverse_mass_matrix = np.ones(1)
step_size = 1e-3
nuts = blackjax.nuts(numba_logpdf, step_size, inverse_mass_matrix)
init = nuts.init(0.)

rng_key, init_key = jax.random.split(rng_key)
state, info = nuts.step(init_key, init)

for _ in range(10):
    rng_key, nuts_key = jax.random.split(rng_key)
    state, _ = nuts.step(nuts_key, state)

print(state)
```

--------------------------------

### Demonstrate JAX incompatibility with Numba functions

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_other_frameworks.md

Attempting to JIT-compile or differentiate Numba-compiled functions directly with JAX will raise exceptions.

```python
import jax.numpy as jnp

try:
    jax.jit(logdensity_fn)(1.)
except Exception:
    print("JAX raised an exception while jit-compiling!")

try:
    jax.grad(logdensity_fn)(1.)
except Exception:
    print("JAX raised an exception while differentiating!")
```

--------------------------------

### Run Inference and Plot Posterior Samples

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_use_pytensor.md

Execute the inference loop and then plot the posterior samples using ArviZ. This step visualizes the results of the MCMC sampling.

```python
states, infos = inference_loop(sample_key, kernel, state, 1000)

import numpy as np
import arviz as az

posterior = {
    "alpha": np.exp(states.position["log_a"]),
    "beta":  np.exp(states.position["log_b"]),
}
idata = az.from_dict(posterior={k: v[None, ...] for k, v in posterior.items()})
az.plot_trace(idata);
```

--------------------------------

### Attempting to differentiate with jax.grad

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/howto_custom_gradients.md

Demonstrates that directly applying `jax.grad` to `f` fails because JAX cannot differentiate through the `minimize` function's internal loops.

```python
# We only want the gradient with respect to `x`
try:
    jax.grad(f, has_aux=True)(0.5, 3)
except Exception as e:
    print(e)
```

--------------------------------

### Profile JAX execution with Perfetto

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/examples/speed_up_guide.md

Use the JAX profiler to trace execution and analyze the output in the Perfetto UI.

```python
with jax.profiler.trace("/tmp/jax-trace"):
    result = run(rng_key, initial_state)
    jax.block_until_ready(result)
# open /tmp/jax-trace in https://ui.perfetto.dev
```

--------------------------------

### Generate Sampling API Wrapper

Source: https://github.com/blackjax-devs/blackjax/blob/main/docs/developer/new_algorithm_guide.md

Wraps a module into a GenerateSamplingAPI dataclass. Exposes .init, .build_kernel, and is callable directly.

```python
my_sampler = generate_top_level_api_from(_my_sampler)
```