### Hello World: NUTS Sampler in Blackjax

Source: https://blackjax-devs.github.io/blackjax/index.html

Demonstrates a basic NUTS sampler implementation. Ensure JAX and Blackjax are installed. This example initializes a NUTS kernel, sets an initial state, and iterates through sampling steps.

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

```

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

### Install PyTensor

Source: https://blackjax-devs.github.io/blackjax/examples/howto_use_pytensor.html

Install PyTensor using pip. This is a prerequisite for running the example.

```bash
pip install pytensor
```

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

### Example Usage

Source: https://blackjax-devs.github.io/blackjax/autoapi/blackjax/mcmc/random_walk/index.html

Demonstrates how to use the random walk kernels, including initialization, stepping, and JIT compilation.

```APIDOC
```python
random_walk = blackjax.additive_step_random_walk(logdensity_fn, blackjax.mcmc.random_walk.normal(sigma))
state = random_walk.init(position)
new_state, info = random_walk.step(rng_key, state)

# JIT compilation for better performance
step = jax.jit(random_walk.step)
new_state, info = step(rng_key, state)
```
```

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

### Set up Eight Schools Example Data

Source: https://blackjax-devs.github.io/blackjax/examples/howto_use_pymc.html

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

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

### SMCState Initialization Examples

Source: https://blackjax-devs.github.io/blackjax/autoapi/blackjax/smc/base/index.html

Demonstrates the structure of particles for SMCState with different posterior distributions. Each example shows how particles can be represented as arrays for univariate, bivariate, or multiple variables.

```python
[ Array([[1.], [1.2], [3.4]]) ]
```

```python
[ Array([[1,2], [3,4], [5,6]]) ]
```

```python
[ Array([[1.], [1.2], [3.4]]), Array([[50.], [51], [55]]) ]
```

```python
[ Array([[1., 2.], [1.2, 0.5], [3.4, 50]]), Array([[50., 51., 52., 51], [51., 52., 52. ,54.], [55., 60, 60, 70]]) ]
```

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

### Import necessary libraries

Source: https://blackjax-devs.github.io/blackjax/examples/howto_other_frameworks.html

Import JAX, NumPy, and other required libraries for the example. Set up a JAX random number generator key.

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

```

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

### Low Rank Adaptation Warmup Example

Source: https://blackjax-devs.github.io/blackjax/autoapi/blackjax/adaptation/low_rank_adaptation/index.html

Demonstrates how to use the `run` method of the `AdaptationAlgorithm` returned by `low_rank_window_adaptation`. After warmup, it shows how to initialize a NUTS kernel with the adapted parameters.

```python
(state, params), info = warmup.run(rng_key, position, num_steps=1000)
nuts = blackjax.nuts(logdensity_fn, **params)
```

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

### Example Usage of blackjax.laplace_dhmc

Source: https://blackjax-devs.github.io/blackjax/autoapi/blackjax/mcmc/laplace_dynamic_hmc/index.html

Example of how to initialize and use the blackjax.laplace_dhmc sampler with specified parameters.

```python
sampler = blackjax.laplace_dhmc(
    log_joint, theta_init=jnp.zeros(n_latent),
    step_size=0.1, inverse_mass_matrix=jnp.ones(d_phi),
    maxiter=100,
)
```

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

### HMC Step Example

Source: https://blackjax-devs.github.io/blackjax/autoapi/blackjax/mcmc/hmc/index.html

Demonstrates a single step of the HMC algorithm using a JIT-compiled kernel. Requires an initialized state and a random number generator key.

```python
step = jax.jit(hmc.step)
new_state, info = step(rng_key, state)
```

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

### Install NumPyro and Funsor

Source: https://blackjax-devs.github.io/blackjax/examples/howto_use_funsor.html

Install the necessary libraries, NumPyro and Funsor, to use Funsor's features with NumPyro models.

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

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

### Additive Step Random Walk Example

Source: https://blackjax-devs.github.io/blackjax/_modules/blackjax/mcmc/random_walk.html

Demonstrates the usage of the `additive_step_random_walk` kernel. It shows how to initialize the chain and perform a single step.

```python
random_walk = blackjax.additive_step_random_walk(logdensity_fn, blackjax.mcmc.random_walk.normal(sigma))
state = random_walk.init(position)
new_state, info = random_walk.step(rng_key, state)
```

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

### Initialize SGNHT kernel and state

Source: https://blackjax-devs.github.io/blackjax/autoapi/blackjax/sgmcmc/sgnht/index.html

This example shows how to initialize a SGNHT kernel using a gradient estimator and then initialize the state for the sampler. Ensure you have a gradient estimator function, a JAX random number key, and an initial position.

```python
grad_estimator = blackjax.sgmcmc.gradients.grad_estimator(logprior_fn, loglikelihood_fn, data_size)
sgnht = blackjax.sgnht(grad_estimator)
state = sgnht.init(rng_key, position)
```

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

### init

Source: https://blackjax-devs.github.io/blackjax/autoapi/blackjax/mcmc/laplace_hmc/index.html

Creates an initial `LaplaceHMCState` by running L-BFGS from a cold start to find `theta*(position)`, then evaluating the Laplace log-marginal and its gradient.

```APIDOC
init(_position : blackjax.types.ArrayLikeTree_, _laplace : blackjax.mcmc.laplace_marginal.LaplaceMarginal_) -> LaplaceHMCState
    """Create an initial `LaplaceHMCState`."""
    
    Parameters:
        position: Initial hyperparameter value `phi`.
        laplace: A `LaplaceMarginal` instance returned by `laplace_marginal_factory()`.
```

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

### Setup and Imports for Funsor and BlackJax

Source: https://blackjax-devs.github.io/blackjax/examples/howto_use_funsor.html

Import necessary libraries including JAX, NumPy, BlackJax, and Funsor. Set the Funsor backend to JAX and import specific Funsor components like Tensor, Variable, and distributions.

```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 1D Gaussian model and sampling parameters

Source: https://blackjax-devs.github.io/blackjax/examples/diagnostics.html

Defines a log-density function for a 1D Gaussian model and sets up parameters for NUTS sampling, including number of chains, samples, step size, and inverse mass matrix.

```python
def logdensity_fn(x):
    return jnp.sum(stats.norm.logpdf(x, 0, 1))

# Sampling parameters
num_chains = 4
num_samples = 1000
step_size = 0.5
inverse_mass_matrix = jnp.ones(1)

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

# Initialize multiple chains
rng_key = jax.random.key(0)
initial_positions = jax.random.normal(rng_key, (num_chains, 1)) * 5 # Disperse starting points
initial_states = jax.vmap(nuts.init)(initial_positions)

# Inference loop using jax.lax.scan
def inference_loop(rng_key, initial_state):
    @jax.jit
    def one_step(state, rng_key):
        state, info = nuts.step(rng_key, state)
        return state, (state, info)

    keys = jax.random.split(rng_key, num_samples)
    _, (states, infos) = jax.lax.scan(one_step, initial_state, keys)
    return states, infos

# Run chains in parallel using vmap
rng_key, sample_key = jax.random.split(rng_key)
sample_keys = jax.random.split(sample_key, num_chains)
states, infos = jax.vmap(inference_loop)(sample_keys, initial_states)

# states.position has shape (num_chains, num_samples, 1)
print(f"Samples shape: {states.position.shape}")
```

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

### Initialize Multiple Chains with jax.vmap

Source: https://blackjax-devs.github.io/blackjax/examples/howto_sample_multiple_chains.html

Prepares initial states for multiple chains by vectorizing the `init` function using `jax.vmap`. This ensures that each chain starts with its own distinct initial parameters.

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

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

### Low Rank Adaptation Algorithm Setup

Source: https://blackjax-devs.github.io/blackjax/autoapi/blackjax/adaptation/low_rank_adaptation/index.html

Sets up an AdaptationAlgorithm for step size and low-rank mass matrix adaptation using a Stan-style warmup schedule. This function is suitable for HMC-family samplers and replaces traditional covariance estimation with a low-rank metric. It returns an AdaptationAlgorithm object with a `run` method for warmup.

```python
low_rank_window_adaptation(_algorithm_ , _logdensity_fn : Callable_, _max_rank : int = 10_, _initial_step_size : float = 1.0_, _target_acceptance_rate : float = 0.8_, _gamma : float = 1.0_, _cutoff : float = 2.0_, _progress_bar : bool = False_, _adaptation_info_fn : Callable = return_all_adapt_info_, _integrator =mcmc.integrators.velocity_verlet_, _** extra_parameters_)
```

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

### Install Blackjax using Conda

Source: https://blackjax-devs.github.io/blackjax/index.html

Use this command to install Blackjax from the conda-forge channel.

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

```

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

### init()

Source: https://blackjax-devs.github.io/blackjax/autoapi/blackjax/sgmcmc/sgld/index.html

Initializes the SGLD kernel state.

```APIDOC
## init()

### Description
Initializes the SGLD kernel state.

### Parameters
None

### Returns
- `blackjax.types.ArrayLikeTree`: The initial state of the SGLD kernel.
```

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

### Install Funsor

Source: https://blackjax-devs.github.io/blackjax/examples/howto_use_funsor.html

Install Funsor using pip. Ensure you are using a version compatible with BlackJax.

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

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

### GenerateVariationalAPI Example

Source: https://blackjax-devs.github.io/blackjax/developer/new_algorithm_guide.html

Example of how to use GenerateVariationalAPI to create a top-level API for a variational inference algorithm.

```python
my_vi = GenerateVariationalAPI(
    _my_vi.as_top_level_api,
    _my_vi.init,
    _my_vi.step,
    _my_vi.sample,
)
```

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

### Install Blackjax using pip

Source: https://blackjax-devs.github.io/blackjax/index.html

Use this command to install the latest version of Blackjax via pip.

```bash
pip install blackjax

```

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

### init()

Source: https://blackjax-devs.github.io/blackjax/autoapi/blackjax/mcmc/laplace_dynamic_hmc/index.html

Creates an initial state for the Laplace Dynamic HMC sampler.

```APIDOC
## Functions
`init`(→ LaplaceDynamicHMCState) | Create an initial `LaplaceDynamicHMCState`.
```

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

### Initialize and Use MCLMC Kernel

Source: https://blackjax-devs.github.io/blackjax/autoapi/blackjax/mcmc/mclmc/index.html

Demonstrates how to initialize and use the MCLMC kernel. Kernels are not JIT-compiled by default and require manual compilation for performance.

```python
mclmc = blackjax.mcmc.mclmc.mclmc(
    logdensity_fn=logdensity_fn,
    L=L,
    step_size=step_size
)
state = mclmc.init(position)
new_state, info = mclmc.step(rng_key, state)
```

```python
step = jax.jit(mclmc.step)
new_state, info = step(rng_key, state)
```

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

### Get JAXified Log-Probability Function

Source: https://blackjax-devs.github.io/blackjax/examples/howto_use_pymc.html

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

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

### init

Source: https://blackjax-devs.github.io/blackjax/_modules/blackjax/mcmc/laplace_dynamic_hmc.html

Create an initial LaplaceDynamicHMCState.

```APIDOC
## init

### Description
Create an initial :class:`LaplaceDynamicHMCState`.

### Parameters
- **position** (ArrayLikeTree) - Initial hyperparameter value ``phi``.
- **laplace** (LaplaceMarginal) - A :class:`~blackjax.mcmc.laplace_marginal.LaplaceMarginal` instance.
- **random_generator_arg** (Array) - Initial value for the quasi-random step-count generator (e.g. a PRNG key or Halton index). When called via the top-level API this is seeded automatically from the ``rng_key`` passed to ``.init``.

### Returns
- **LaplaceDynamicHMCState** - An initial state for the Laplace dynamic HMC sampler.
```

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

### Initialize and Use HMC Kernel

Source: https://blackjax-devs.github.io/blackjax/_modules/blackjax/mcmc/hmc.html

Demonstrates the basic initialization and usage of an HMC kernel. Kernels are not JIT-compiled by default and require manual compilation for performance.

```python
hmc = blackjax.hmc(
    logdensity_fn, step_size, inverse_mass_matrix, num_integration_steps
)
state = hmc.init(position)
new_state, info = hmc.step(rng_key, state)
```

```python
step = jax.jit(hmc.step)
new_state, info = step(rng_key, state)
```

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

### Define Treatment Effects Data

Source: https://blackjax-devs.github.io/blackjax/examples/howto_use_tfp.html

Sets up the observed treatment effects and their standard deviations for the Eight Schools example.

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

```

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

### Regression Test Case Registration

Source: https://blackjax-devs.github.io/blackjax/developer/new_algorithm_guide.html

Example of how to add a new algorithm to the accuracy regression suite by defining its parameters and sampling configuration.

```python
{
    "algorithm": blackjax.my_sampler,
    "initial_position": {"log_scale": 0.0, "coefs": 4.0},
    "parameters": {"step_size": 0.1},
    "num_warmup_steps": 1_000,
    "num_sampling_steps": 3_000,
},
```

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

### Initialize and Use Additive Step Random Walk Kernel

Source: https://blackjax-devs.github.io/blackjax/_modules/blackjax/mcmc/random_walk.html

Demonstrates how to initialize and use the `additive_step_random_walk` kernel. This is useful for creating custom random walk proposals.

```python
rw = blackjax.additive_step_random_walk(logdensity_fn, random_step)
state = rw.init(position)
new_state, info = rw.step(rng_key, state)
```

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

### init()

Source: https://blackjax-devs.github.io/blackjax/autoapi/blackjax/mcmc/ghmc/index.html

Initializes the GHMC sampler state.

```APIDOC
def init() -> GHMCState:
    """Initialize the GHMCState."""
    ...
```

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

### init

Source: https://blackjax-devs.github.io/blackjax/autoapi/blackjax/smc/inner_kernel_tuning/index.html

Initialize the inner-kernel-tuning SMC state.

```APIDOC
def init(alg_init_fn, position, initial_parameter_value):
    """Initialize the inner-kernel-tuning SMC state."""
    pass
```

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

### init

Source: https://blackjax-devs.github.io/blackjax/autoapi/blackjax/mcmc/laplace_hmc/index.html

Create an initial `LaplaceHMCState` for the Laplace-HMC sampler.

```APIDOC
def init(position: Any, logdensity: Any, logdensity_grad: Any, theta_star: Any) -> LaplaceHMCState
```

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

### Initialize Tempered SMC State

Source: https://blackjax-devs.github.io/blackjax/_modules/blackjax/smc/tempered.html

Initializes the state for the Tempered SMC algorithm. Use this when starting a new Tempered SMC chain.

```python
state = smc.tempered.init(initial_particles)
```

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

### Visualize Posterior Samples with ArviZ

Source: https://blackjax-devs.github.io/blackjax/examples/howto_use_pymc.html

Generates a trace plot of the posterior samples using ArviZ and Matplotlib. Requires `matplotlib` and `arviz` to be installed.

```python
import matplotlib.pyplot as plt
import arviz as az

idata = az.from_dict(
    posterior={k: v[None, ...]
               for k, v in zip(model.initial_point().keys(), states.position)})
az.plot_trace(idata)
plt.tight_layout()
```

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

### Initialize NUTS State

Source: https://blackjax-devs.github.io/blackjax/examples/quickstart.html

Initialize the NUTS sampler with the starting position for the MCMC chain. The initial state includes the position and the log-density at that position.

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

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

### Initialize and Step RMH Kernel

Source: https://blackjax-devs.github.io/blackjax/autoapi/blackjax/mcmc/random_walk/index.html

Demonstrates how to initialize an RMH kernel with a log-density function and proposal generator, and then perform a single step.

```python
rmh = blackjax.rmh(logdensity_fn, proposal_generator)
state = rmh.init(position)
new_state, info = rmh.step(rng_key, state)
```

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

### Initializing LaplaceHMCState

Source: https://blackjax-devs.github.io/blackjax/_modules/blackjax/mcmc/laplace_hmc.html

Creates an initial `LaplaceHMCState` by running L-BFGS from a cold start to find `theta_star(position)`, and then evaluating the Laplace log-marginal and its gradient.

```python
def init(
    position: ArrayLikeTree,
    laplace: LaplaceMarginal,
) -> LaplaceHMCState:
    """Create an initial :class:`LaplaceHMCState`.

    Runs L-BFGS from cold start to find ``theta*(position)``, then evaluates
    the Laplace log-marginal and its gradient.

    Parameters
    ----------
    position
        Initial hyperparameter value ``phi``.
    laplace
        A :class:`~blackjax.mcmc.laplace_marginal.LaplaceMarginal` instance
        returned by :func:`~blackjax.mcmc.laplace_marginal.laplace_marginal_factory`.
    """
    (logdensity, theta_star), logdensity_grad = jax.value_and_grad(
        laplace, has_aux=True
    )(position)
    return LaplaceHMCState(position, logdensity, logdensity_grad, theta_star)
```

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

### init()

Source: https://blackjax-devs.github.io/blackjax/autoapi/blackjax/smc/tempered/index.html

Initializes the Tempered SMC state.

```APIDOC
def init() -> TemperedSMCState:
    """Initialize the Tempered SMC state."""
    ...
```

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

### Import Necessary Libraries

Source: https://blackjax-devs.github.io/blackjax/examples/howto_use_funsor.html

Import NumPyro, its distributions, Funsor's configuration decorator, and initialization utilities.

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

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

### Example output of inference statistics

Source: https://blackjax-devs.github.io/blackjax/examples/howto_laplace_hmc.html

Displays the expected output format for the inference statistics, including acceptance rate, divergences, and the posterior mean of phi.

```text
Acceptance rate:  0.80
Divergences:      0
phi posterior mean:  0.865  (true log_sigma=0.693)

```

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

### init()

Source: https://blackjax-devs.github.io/blackjax/autoapi/blackjax/mcmc/laplace_dynamic_hmc/index.html

Initializes the Laplace Dynamic HMC sampler state.

```APIDOC
## init()

### Description
Initializes the Laplace Dynamic HMC sampler state.

### Parameters
- **phi_init** (object): Initial value for the latent variable.
- **rng_key** (object): JAX random number generator key.

### Returns
- **state** (LaplaceDynamicHMCState): The initial state of the sampler.
```

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

### Get Final Step Size and Mass Matrix

Source: https://blackjax-devs.github.io/blackjax/_modules/blackjax/adaptation/window_adaptation.html

Returns the final adapted step size and inverse mass matrix after the adaptation process is complete.

```python
def final(warmup_state: WindowAdaptationState) -> tuple[float, Array]:
        """Return the final values for the step size and mass matrix."""
        step_size = jnp.exp(warmup_state.ss_state.log_step_size_avg)
        inverse_mass_matrix = warmup_state.imm_state.inverse_mass_matrix
        return step_size, inverse_mass_matrix
```

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

### Run NUTS with Adapted Parameters

Source: https://blackjax-devs.github.io/blackjax/examples/quickstart.html

Initialize and run the NUTS kernel using parameters obtained from window adaptation. This allows for efficient sampling without manual tuning.

```python
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()
```

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

### Load Iris Dataset with Scikit-learn

Source: https://blackjax-devs.github.io/blackjax/examples/howto_use_oryx.html

Loads the iris dataset using scikit-learn for use in Blackjax and Oryx examples. Ensures features and labels are correctly extracted.

```python
from sklearn import datasets

iris = datasets.load_iris()
features, labels = iris['data'], iris['target']
num_features = features.shape[-1]
num_classes = len(iris.target_names)
```

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

### init

Source: https://blackjax-devs.github.io/blackjax/autoapi/blackjax/sgmcmc/sgnht/index.html

Initializes the SGNHT state.

```APIDOC
init(_position : blackjax.types.ArrayLikeTree_, _rng_key : blackjax.types.PRNGKey_, _xi : float_) -> SGNHTState
```

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

### Example Output of Array Shapes

Source: https://blackjax-devs.github.io/blackjax/examples/howto_sample_multiple_chains.html

This output shows the shape comparison between the original and sharded arrays, illustrating how the data has been distributed across devices for parallel processing.

```text
((2000, 4), (4, 2000))

```

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

### init

Source: https://blackjax-devs.github.io/blackjax/_modules/blackjax/mcmc/laplace_hmc.html

Create an initial LaplaceHMCState.

```APIDOC
## init

### Description
Create an initial :class:`LaplaceHMCState`.

Runs L-BFGS from cold start to find ``theta*(position)``, then evaluates the Laplace log-marginal and its gradient.

### Parameters
- **position** (ArrayLikeTree) - Initial hyperparameter value ``phi``.
- **laplace** (LaplaceMarginal) - A :class:`~blackjax.mcmc.laplace_marginal.LaplaceMarginal` instance returned by :func:`~blackjax.mcmc.laplace_marginal.laplace_marginal_factory`.
```

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

### init

Source: https://blackjax-devs.github.io/blackjax/_modules/blackjax/mcmc/adjusted_mclmc.html

Create an initial state for the MHMCHMC kernel.

```APIDOC
## init

### Description
Create an initial state for the MHMCHMC kernel.

### Parameters
* **position** (ArrayLikeTree) - Initial position of the chain.
* **logdensity_fn** (Callable) - Log-density function of the target distribution.

### Returns
* The initial HMCState.
```

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

### init

Source: https://blackjax-devs.github.io/blackjax/_modules/blackjax/adaptation/meads_adaptation.html

Initializes the MEADS adaptation state.

```APIDOC
## init

### Description
Initialize with parameters computed from all chains, replicated per fold.

### Parameters

*   **positions** (ArrayLikeTree) - Initial positions of the chains.
*   **logdensity_grad** (ArrayLikeTree) - Gradients of the log density function at the initial positions.
*   **num_folds** (int) - The number of folds to divide the chains into.

### Returns

*   **MEADSAdaptationState** - The initialized adaptation state.
```

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

### Initialize JAX and RNG Key

Source: https://blackjax-devs.github.io/blackjax/examples/howto_use_oryx.html

Initializes JAX and sets up a random number generator key using the current date. This is a common setup for JAX-based probabilistic programming.

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

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

### Numpydoc Docstring Example

Source: https://blackjax-devs.github.io/blackjax/developer/design_principles.html

Public functions, classes, and modules must have Numpydoc-formatted docstrings. Include Parameters and Returns sections. Explain magic numbers and mathematical derivations inline.

```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)``.
    """

```

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

### JIT compile SGNHT step function

Source: https://blackjax-devs.github.io/blackjax/autoapi/blackjax/sgmcmc/sgnht/index.html

To optimize performance, it is recommended to JIT compile the SGNHT step function. This example shows how to apply `jax.jit` to the `sgnht.step` function before using it.

```python
step = jax.jit(sgnht.step)
new_state = step(rng_key, state, minibatch, step_size)
```

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

### MEADS Adaptation Algorithm

Source: https://blackjax-devs.github.io/blackjax/_modules/blackjax/adaptation/meads_adaptation.html

The `AdaptationAlgorithm` class facilitates the MEADS adaptation process. The `run` method initializes the adaptation, performs a specified number of steps, and returns the final adapted parameters and adaptation information.

```APIDOC
## AdaptationAlgorithm

### Description

An adaptation algorithm that adapts the step size and momentum parameters of a kernel.

### Parameters

- **run** (Callable) - A function that takes a PRNGKey, initial positions, and the number of steps, and returns the adapted states and parameters.

### Returns

- An instance of `AdaptationAlgorithm`.
```

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

### buildtree_integrate

Source: https://blackjax-devs.github.io/blackjax/_modules/blackjax/mcmc/trajectory.html

Integrates the Hamiltonian trajectory starting from an initial state using a recursive tree doubling scheme until a termination criterion is met. This function is central to building proposals in HMC.

```APIDOC
## buildtree_integrate

### Description
Integrate the trajectory starting from `initial_state` and update the proposal recursively with tree doubling until the termination criterion is met. The function `buildtree_integrate` calls itself for tree_depth > 0, thus invokes the recursive scheme that builds a trajectory by doubling a binary tree.

### Parameters
#### Parameters
- **integrator**: The symplectic integrator used to integrate the hamiltonian trajectory.
- **kinetic_energy**: Function to compute the current value of the kinetic energy.
- **uturn_check_fn**: Determines whether the termination criterion has been met.
- **divergence_threshold**: Value of the difference of energy between two consecutive states above which we say a transition is divergent.
- **use_robust_uturn_check**: Bool to indicate whether to perform additional U turn check between two trajectory.

#### Parameters for buildtree_integrate function:
- **rng_key**: Key used by JAX's random number generator.
- **initial_state**: The initial state from which we start expanding the trajectory.
- **direction** (int in {-1, 1}): The direction in which to expand the trajectory.
- **tree_depth**: The depth of the binary tree doubling.
- **step_size**: The step size of the symplectic integrator.
- **initial_energy**: Initial energy H0 of the HMC step (not to confused with the initial energy of the subtree)

### Returns
A ``buildtree_integrate(rng_key, state, direction, tree_depth, step_size, initial_energy)`` function returning ``(rng_key, proposal, trajectory, is_diverging, is_turning)``.
```

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

### init

Source: https://blackjax-devs.github.io/blackjax/_modules/blackjax/smc/partial_posteriors_path.html

Initializes the state for the Partial Posteriors SMC algorithm.

```APIDOC
## init

### Description
Initializes the state for the Partial Posteriors SMC algorithm. The initial `data_mask` is all zeros, meaning no likelihood term is added initially, only the prior.

### Parameters
- **particles** (ArrayLikeTree) - The initial particles' positions.
- **num_datapoints** (int) - The total number of observations that could potentially be used in a partial posterior.

### Returns
- **PartialPosteriorsSMCState** - The initial state of the SMC algorithm.
```

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

### Initialize Pretuning SMC State

Source: https://blackjax-devs.github.io/blackjax/_modules/blackjax/smc/pretuning.html

Initializes a state for pretuning SMC, wrapping the initial SMC state and MCMC parameters. Use this when starting a sampling process that involves pretuning.

```python
def init(alg_init_fn, position, initial_parameter_value):
    """Initialize a pretuning SMC state.

    Parameters
    ----------
    alg_init_fn
        The ``init`` function of the underlying SMC algorithm.
    position
        Initial particle positions (PyTree).
    initial_parameter_value
        Initial dict of MCMC parameters assigned to each chain.

    Returns
    -------
    A StateWithParameterOverride wrapping the SMC state and the initial
    parameter distribution.
    """
    return StateWithParameterOverride(alg_init_fn(position), initial_parameter_value)
```

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

### Initialize Random Walk State

Source: https://blackjax-devs.github.io/blackjax/_modules/blackjax/mcmc/random_walk.html

Creates the initial state for a Random Walk (RW) chain. This function requires the starting position and the log-probability density function of the target distribution.

```python
def init(position: ArrayLikeTree, logdensity_fn: Callable) -> RWState:
    """Create a chain state from a position.

    Parameters
    ----------
    position
        The initial position of the chain.
    logdensity_fn
        Log-probability density function of the distribution we wish to sample
        from.

    Returns
    -------
    The initial state of the RW chain.
    """
    return RWState(position, logdensity_fn(position))
```

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

### Compute Summary Statistics for Adaptation

Source: https://blackjax-devs.github.io/blackjax/_modules/blackjax/adaptation/laps_burn_in.html

Calculates various statistics from the current state and trajectory information to guide the adaptation process. Includes equipartition terms, position statistics, and energy changes.

```python
def summary_statistics_fn(self, state, info, rng_key):
        position_flat, unravel_fn = ravel_pytree(state.position)

        return {
            "equipartition_diagonal": equipartition_diagonal(state),
            "equipartition_fullrank": equipartition_fullrank(state, rng_key),
            "x": position_flat,
            "xsq": jnp.square(position_flat),
            "E": info["energy_change"],
            "Esq": jnp.square(info["energy_change"]),
            "rejection_rate_nans": info["nans"],
            "observables_for_bias": self.observables_for_bias(state.position),
            "observables": self.observables(state.position),
            "entropy": -info["logdensity"],
            # "uturn": jnp.sqrt(jnp.sum(jnp.square(state.logdensity_grad - jnp.dot(state.logdensity_grad, state.momentum) * state.momentum))) / (self.ndims - 1)
        }
```

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

### build_pretune() Function

Source: https://blackjax-devs.github.io/blackjax/autoapi/blackjax/smc/pretuning/index.html

Implements the pretuning procedure described by Buchholz et al. (https://arxiv.org/pdf/1808.07730).

```APIDOC
def build_pretune(mcmc_init_fn: Callable, mcmc_step_fn: Callable, alpha: float, ...) -> Any:
    """Implements Buchholz et al https://arxiv.org/pdf/1808.07730 pretuning procedure."""
    ...
```