### FlowJAX Simple Distribution Example: Normal

Source: https://github.com/danielward27/flowjax/blob/main/docs/getting_started.rst

Demonstrates the behavior of FlowJAX distributions using the Normal distribution. Shows how to initialize a distribution, sample from it (single and batched), and evaluate log probabilities.

```Python
from flowjax.distributions import Normal
import jax.numpy as jnp
import jax.random as jr

# Initialize a Normal distribution
normal = Normal(loc=jnp.arange(3), scale=1)
print(normal.shape)

# Sample from the distribution
key = jr.key(0)
sample = normal.sample(key)
print(sample.shape)
batch = normal.sample(key, (4,))
print(batch.shape)

# Evaluate log probabilities
log_prob = normal.log_prob(sample)
print(log_prob.shape)
log_probs = normal.log_prob(batch)
print(log_probs.shape)
```

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

### FlowJAX Conditional Distribution Example

Source: https://github.com/danielward27/flowjax/blob/main/docs/getting_started.rst

Illustrates the use of conditional distributions in FlowJAX. Shows how to create a conditional distribution using `coupling_flow`, inspect its shape and conditional shape, and sample/calculate log probabilities with conditioning variables.

```Python
from flowjax.flows import coupling_flow
from flowjax.distributions import Normal
import jax.numpy as jnp
import jax.random as jr

# Create a conditional distribution
key = jr.key(0)
dist = coupling_flow(key, base_dist=Normal(jnp.zeros(3)), cond_dim=2)
print(dist.shape)
print(dist.cond_shape)

# Sampling with a single conditioning variable instance
condition = jnp.ones(2)
samples = dist.sample(key, (10,), condition=condition)
print(samples.shape)
print(dist.log_prob(samples, condition).shape)

# Sampling with multiple conditioning variable instances
condition = jnp.ones((5, 2))
samples = dist.sample(key, condition=condition)
print(samples.shape)
print(dist.log_prob(samples, condition).shape)
```

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

### FlowJAX Simple Bijection Example: Affine Transformation

Source: https://github.com/danielward27/flowjax/blob/main/docs/getting_started.rst

Demonstrates a simple bijection using the Affine transformation in FlowJAX. Shows how to perform forward and inverse transformations, and compute the log absolute value of the Jacobian determinant.

```Python
import jax.numpy as jnp
from flowjax.bijections import Affine

# Initialize an Affine bijection
bijection = Affine(loc=0, scale=2)
x = 1

# Forward transformation
y = bijection.transform(x)
print(y)

# Inverse transformation
print(bijection.inverse(y))

# Transformation and log determinant
print(bijection.transform_and_log_det(x))
print(bijection.inverse_and_log_det(y))
```

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

### FlowJAX Transformed Distribution: Log-Normal

Source: https://github.com/danielward27/flowjax/blob/main/docs/getting_started.rst

Illustrates creating a log-normal distribution using FlowJAX's `Transformed` class, combining a base distribution with a transformation (Exp bijection).

```Python
from flowjax.distributions import Normal, Transformed
from flowjax.bijections import Exp

# Create a log-normal distribution
log_normal = Transformed(Normal(), Exp())
```

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

### FlowJAX Bijection Vectorization with JAX vmap

Source: https://github.com/danielward27/flowjax/blob/main/docs/getting_started.rst

Shows how to vectorize bijection methods using JAX's `vmap` for applying transformations to arrays element-wise.

```Python
import jax
import jax.numpy as jnp
from flowjax.bijections import Scale

# Initialize a Scale bijection
scale = Scale(2)  # shape ()
x = jnp.arange(3)

# Vectorize the transform method
vectorized_transform = jax.vmap(scale.transform)
print(vectorized_transform(x))
```

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

### Vectorize Normal Distribution Initialization with Equinox

Source: https://github.com/danielward27/flowjax/blob/main/docs/getting_started.rst

Shows how to create a batch of independent Normal distributions using `equinox.filter_vmap`. This is useful for defining multiple distributions with potentially different parameters.

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

normals = eqx.filter_vmap(Normal)(jnp.arange(3))  # batch of normals with shape ()

```

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

### Vectorize Log Probability Calculation with Equinox

Source: https://github.com/danielward27/flowjax/blob/main/docs/getting_started.rst

Illustrates how to compute the log probability for a batch of distributions using `equinox.filter_vmap`. This allows efficient calculation across multiple distribution instances and data points.

```python
log_probs = eqx.filter_vmap(lambda dist, x: dist.log_prob(x))(normals, jnp.arange(3))

```

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

### Define LogNormal Distribution using Equinox

Source: https://github.com/danielward27/flowjax/blob/main/docs/getting_started.rst

Demonstrates the definition of a LogNormal distribution by composing a Normal distribution with an Exponential bijection. It highlights the requirement for matching shapes between the bijection and the base distribution.

```python
class LogNormal(AbstractTransformed):
    base_dist: Normal
    bijection: Exp

    def __init__(self, loc: ArrayLike = 0, scale: ArrayLike = 1):
        self.base_dist = Normal(loc, scale)
        self.bijection = Exp(self.base_dist.shape)

```

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

### Install FlowJAX for Development

Source: https://github.com/danielward27/flowjax/blob/main/README.md

This sequence of commands clones the FlowJAX repository, installs it in editable mode with development dependencies, and installs Pandoc, which is required for building the documentation.

```Bash
git clone https://github.com/danielward27/flowjax.git
cd flowjax
pip install -e .[dev]
sudo apt-get install pandoc  # Required for building documentation

```

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

### Install FlowJAX using pip

Source: https://github.com/danielward27/flowjax/blob/main/docs/index.rst

This command installs the FlowJAX package using pip. Ensure you have Python and pip installed.

```bash
pip install flowjax
```

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

### Install FlowJAX

Source: https://github.com/danielward27/flowjax/blob/main/README.md

This command installs the FlowJAX package using pip, making it available for use in Python projects.

```Bash
pip install flowjax

```

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

### JIT Compilation for Sampling Batches

Source: https://github.com/danielward27/flowjax/blob/main/docs/faq.rst

Explains the benefit of using JIT compilation for performance, particularly when sampling multiple batches. The example contrasts slow, non-jitted sampling with fast, jitted sampling using `eqx.filter_jit`.

```Python
import equinox as eqx
import jax.random as jr

batch_size = 256
keys = jr.split(jr.key(0), 5)

# Fast - sample jitted!
results = []
for batch_key in keys:
    x = eqx.filter_jit(flow.sample)(batch_key, (batch_size,)) # Assuming 'flow' is defined
    results.append(x)
```

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

### Standardize Variables with Affine Bijection

Source: https://github.com/danielward27/flowjax/blob/main/docs/faq.rst

Illustrates standardizing data using an `Affine` bijection for preprocessing. The example shows transforming data, fitting a FlowJAX model to the processed data, and then inverting the preprocessing transformation to 'undo' it.

```Python
import jax
from flowjax.bijections import Affine, Invert
from flowjax.distributions import Transformed
from flowjax.train import fit_to_data
import jax.numpy as jnp
import jax.random as jr

key = jr.key(0)
x = jr.normal(key, (1000,3))
flow = Normal(jnp.ones(3))

preprocess = Affine(-x.mean(axis=0)/x.std(axis=0), 1/x.std(axis=0))
x_processed = jax.vmap(preprocess.transform)(x)
flow, losses = fit_to_data(key, dist=flow, x=x_processed) # doctest: +SKIP
flow = Transformed(flow, Invert(preprocess))
```

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

### Test Flowjax Documentation Doctests

Source: https://github.com/danielward27/flowjax/blob/main/docs/README.md

Tests the doctest code blocks within the Flowjax documentation. Run this command from the top-level directory of the project to ensure all embedded code examples are valid and functional.

```bash
make -C docs doctest
```

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

### Create and Train a Normalizing Flow Model with FlowJAX

Source: https://github.com/danielward27/flowjax/blob/main/README.md

This example demonstrates how to create a Block Neural Autoregressive Flow model, train it on toy data using maximum likelihood, and then evaluate its log-probability and sample from the trained distribution. It utilizes JAX for numerical operations and random number generation, and FlowJAX for flow components and training utilities.

```Python
from flowjax.flows import block_neural_autoregressive_flow
from flowjax.train import fit_to_data
from flowjax.distributions import Normal
import jax.random as jr
import jax.numpy as jnp

data_key, flow_key, train_key, sample_key = jr.split(jr.key(0), 4)

x = jr.uniform(data_key, (5000, 2))  # Toy data

flow = block_neural_autoregressive_flow(
    key=flow_key,
    base_dist=Normal(jnp.zeros(x.shape[1])),
)

flow, losses = fit_to_data(
    key=train_key,
    dist=flow,
    data=x,
    learning_rate=5e-3,
    max_epochs=200,
    )

# We can now evaluate the log-probability of arbitrary points
log_probs = flow.log_prob(x)

# And sample the distribution
samples = flow.sample(sample_key, (1000, ))

```

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

### Runtime Type Checking with Jaxtyping and Beartype

Source: https://github.com/danielward27/flowjax/blob/main/docs/faq.rst

Demonstrates how to enable runtime type checking in FlowJAX using `jaxtyping`'s import hook and `beartype`. This setup provides helpful error messages for type mismatches, as shown in the example where an incorrect shape is provided to `Exp`.

```Python
from jaxtyping import install_import_hook

with install_import_hook("flowjax", "beartype.beartype"):
   from flowjax import bijections as bij

exp = bij.Exp(shape=2) # Raises a helpful error as 2 is not a tuple
```

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

### Create Masked Autoregressive Flow with Spline Transformer

Source: https://github.com/danielward27/flowjax/blob/main/docs/examples/unconditional.ipynb

Initializes a masked autoregressive flow model. It uses a Normal distribution as the base and a RationalQuadraticSpline as the transformer, configured with 8 knots over an interval of 4. This setup is suitable for approximating complex data distributions.

```Python
key, subkey = jr.split(jr.key(0))
flow = masked_autoregressive_flow(
    subkey,
    base_dist=Normal(jnp.zeros(x.shape[1])),
    transformer=RationalQuadraticSpline(knots=8, interval=4),
)
```

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

### Build Flowjax Documentation Locally

Source: https://github.com/danielward27/flowjax/blob/main/docs/README.md

Builds the HTML documentation for the Flowjax project locally using Sphinx. Navigate to the 'docs' directory and execute the 'make html' command. The generated documentation can be accessed at './docs/_build/html/index.html'.

```bash
make -C docs html
```

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

### Import Libraries for FlowJax

Source: https://github.com/danielward27/flowjax/blob/main/docs/examples/conditional.ipynb

Imports necessary libraries from JAX and FlowJax for building and training normalizing flows, including JAX NumPy, random number generation, plotting, and specific flow and training modules.

```Python
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt

from flowjax.distributions import Normal
from flowjax.flows import block_neural_autoregressive_flow
from flowjax.train import fit_to_data
```

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

### Initialize GaussianMixtureSimulator and Generate Observation

Source: https://github.com/danielward27/flowjax/blob/main/docs/examples/snpe.ipynb

Initializes the GaussianMixtureSimulator task and generates a single 'observation' by sampling a true parameter 'theta_true' from the prior and then simulating data based on this parameter. This sets up the data generation process for the SNPE task.

```Python
task = GaussianMixtureSimulator()

# Generating an "observation"
key, subkey = jr.split(jr.key(0))
theta_true = task.prior.sample(subkey)
key, subkey = jr.split(key)
observed = task.simulator(subkey, theta_true)
```

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

### Import Libraries for FlowJax

Source: https://github.com/danielward27/flowjax/blob/main/docs/examples/constrained.ipynb

Imports necessary libraries from JAX, Matplotlib, and FlowJax for building and training normalizing flows. This includes modules for distributions, bijections, and training utilities.

```Python
import jax
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from paramax import non_trainable

import flowjax.bijections as bij
from flowjax.distributions import Normal, Transformed
from flowjax.flows import masked_autoregressive_flow
from flowjax.train import fit_to_data
```

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

### Extract Parameters using Equinox

Source: https://github.com/danielward27/flowjax/blob/main/docs/faq.rst

Shows how to partition parameters from FlowJAX distributions or bijections using `equinox.partition`. This function separates inexact arrays (parameters) from static components, specifically filtering out arrays marked as `paramax.NonTrainable` into the static part.

```Python
import equinox as eqx

params, static = eqx.partition(
    dist,
    eqx.is_inexact_array,
    is_leaf=lambda leaf: isinstance(leaf, paramax.NonTrainable),
)
```

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

### FlowJAX Citation

Source: https://github.com/danielward27/flowjax/blob/main/README.md

This BibTeX entry provides the necessary information to cite the FlowJAX package in academic work. Users should replace '[version number]' and '[release year of version]' with the appropriate details.

```BibTeX
@software{ward2023flowjax,
  title = {FlowJAX: Distributions and Normalizing Flows in Jax},
  author = {Daniel Ward},
  url = {https://github.com/danielward27/flowjax},
  version = {[version number]},
  year = {[release year of version]},
  doi = {10.5281/zenodo.10402073},
}

```

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

### Import Libraries for FlowJax SNPE

Source: https://github.com/danielward27/flowjax/blob/main/docs/examples/snpe.ipynb

Imports necessary libraries from Equinox, JAX, Matplotlib, and FlowJax for implementing Sequential Neural Posterior Estimation. This includes modules for neural flows, distributions, training utilities, and loss functions.

```Python
import equinox as eqx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from jax import vmap

import flowjax.bijections as bij
from flowjax.distributions import Normal, Transformed
from flowjax.flows import masked_autoregressive_flow
from flowjax.tasks import GaussianMixtureSimulator
from flowjax.train import fit_to_data
from flowjax.train.losses import ContrastiveLoss, MaximumLikelihoodLoss
```

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

### FlowJAX Bijection Composition with Chain

Source: https://github.com/danielward27/flowjax/blob/main/docs/api/flows.rst

Illustrates the use of flowjax.bijections.Chain for composing arbitrary and heterogeneous bijections. This method compiles each layer separately.

```Python
from flowjax.bijections import Chain

# Example usage (conceptual):
# chain_bijection = Chain([bijection1, bijection2, ...])
```

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

### NumPyro Regression Model with FlowJAX Sample

Source: https://github.com/danielward27/flowjax/blob/main/docs/api/experimental.rst

Demonstrates a simple regression model using FlowJAX distributions within a NumPyro model. It replaces NumPyro's `sample` with `flowjax.experimental.numpyro.sample` to incorporate FlowJAX distributions, specifically a Normal distribution for the regression coefficients.

```Python
from numpyro.infer import MCMC, NUTS
from flowjax.experimental.numpyro import sample
from flowjax.distributions import Normal
import jax.random as jr
import numpy as np

def numpyro_model(X, y):
    """Example regression model defined in terms of FlowJAX distributions"""
    beta = sample("beta", Normal(np.zeros(2)))
    sample("y", Normal(X @ beta), obs=y)

X = np.random.randn(100, 2)
beta_true = np.array([-1, 1])
y = X @ beta_true + np.random.randn(100)
mcmc = MCMC(NUTS(numpyro_model), num_warmup=10, num_samples=100)
mcmc.run(jr.key(0), X, y)
```

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

### FlowJAX Transformed Distribution

Source: https://github.com/danielward27/flowjax/blob/main/docs/api/flows.rst

Shows how FlowJAX flows are constructors for flowjax.distributions.Transformed distributions. These distributions are built by composing multiple bijections.

```Python
from flowjax.distributions import Transformed

# Example usage (conceptual):
# transformed_dist = Transformed(base_distribution, flow_bijection)
```

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

### Visualize Learned Densities and Constraints (FlowJax, JAX, Matplotlib)

Source: https://github.com/danielward27/flowjax/blob/main/docs/examples/constrained.ipynb

This snippet visualizes the learned densities of different flow configurations, highlighting the impact of constraints. It samples data from the original distribution, an unconstrained flow, and two constrained flows, then plots the samples, coloring valid samples green and invalid samples red.

```Python
key, subkey = jr.split(key)
samples = {
    "True\ndistribution": x,
    "No\nconstraints": flow_no_constraints.sample(subkey, (x.shape[0],)),
    "Constrained\n(option 1)": flow_constrained_1.sample(subkey, (x.shape[0],)),
    "Constrained\n(option 2)": flow_constrained_2.sample(subkey, (x.shape[0],)),
}

fig = plt.figure(figsize=(6, 5))
axes = fig.subplots(ncols=len(samples), sharex=True, sharey=True)

for (k, v), ax in zip(samples.items(), axes, strict=True):
    is_valid = jnp.logical_and(v[:, 0] >= 0, v[:, 0] <= 1)
    ax.scatter(v[:, 0][is_valid], v[:, 1][is_valid], s=0.1, color="tab:green")
    ax.scatter(v[:, 0][~is_valid], v[:, 1][~is_valid], s=0.1, color="tab:red")
    ax.set_title(k)
    ax.set_box_aspect(1)
plt.show()
```

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

### Evaluate Density and Sample from FlowJax Model

Source: https://github.com/danielward27/flowjax/blob/main/docs/examples/unconditional.ipynb

Demonstrates how to use the trained FlowJax model to evaluate the log probability density of arbitrary points and generate new samples. The sampling process is visualized by plotting the true data samples against the samples generated by the flow.

```Python
five_points = jnp.ones((5, 2))
flow.log_prob(five_points)

key, subkey = jr.split(key)
x_samples = flow.sample(subkey, (n_samples,))

fig, axs = plt.subplots(ncols=2)

axs[0].scatter(x[:, 0], x[:, 1], s=0.1)
axs[0].set_title("True samples")

axs[1].scatter(x_samples[:, 0], x_samples[:, 1], s=0.1)
axs[1].set_title("Flow samples")

lims = (-2.5, 2.5)
for ax in axs:
    ax.set_xlim(lims)
    ax.set_ylim(lims)
    ax.set_aspect("equal")

plt.tight_layout()
plt.show()
```

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

### Import Libraries for FlowJax Density Estimation

Source: https://github.com/danielward27/flowjax/blob/main/docs/examples/unconditional.ipynb

Imports necessary libraries from JAX, FlowJax, and Matplotlib for density estimation tasks. This includes JAX for numerical operations and random number generation, FlowJax for flow architectures and distributions, and Matplotlib for plotting.

```Python
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt

from flowjax.bijections import RationalQuadraticSpline
from flowjax.distributions import Normal
from flowjax.flows import masked_autoregressive_flow
from flowjax.tasks import two_moons
from flowjax.train import fit_to_data
```

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

### Fit Flow to Data with Constraints (FlowJax, JAX)

Source: https://github.com/danielward27/flowjax/blob/main/docs/examples/constrained.ipynb

This code fits a specified flow model to the provided data, incorporating parameters for learning rate, patience, and epochs. It's designed to work with both constrained and unconstrained flow configurations within the FlowJax framework.

```Python
flow2, losses = fit_to_data(
    key=subkey,
    dist=unbounded_flow,
    data=x_unbounded,
    learning_rate=5e-3,
    max_patience=10,
    max_epochs=70,
)
```

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

### Transform Data for Unbounded Flow Support (FlowJax, JAX)

Source: https://github.com/danielward27/flowjax/blob/main/docs/examples/constrained.ipynb

This snippet shows how to transform data to match the support requirements of an unbounded flow model using FlowJax and JAX. It utilizes the `inverse` method of a transformation object and `jax.vmap` for efficient mapping across data points.

```Python
x_unbounded = jax.vmap(to_constrained.inverse)(x_clipped)
```

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

### FlowJAX Bijection Composition with Scan

Source: https://github.com/danielward27/flowjax/blob/main/docs/api/flows.rst

Demonstrates the use of flowjax.bijections.Scan for composing bijections with a shared structure, which avoids compiling each layer separately and reduces compilation times. This is the default method used in FlowJAX flows.

```Python
from flowjax.bijections import Scan

# Example usage (conceptual):
# scan_bijection = Scan(bijection_template, num_layers=...)
```

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

### Fit Flow to Data using FlowJAX

Source: https://github.com/danielward27/flowjax/blob/main/docs/api/training.rst

The `fit_to_data` function in FlowJAX allows users to train a flow model by fitting it to samples from a distribution. It can also accommodate corresponding conditioning variables if needed for the specific use case.

```python
from flowjax.train import fit_to_data

# Example usage (assuming you have data and conditioning variables)
# fit_to_data(flow, data, conditioning_variables=None, ...)
```

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

### Control Training Step in FlowJAX

Source: https://github.com/danielward27/flowjax/blob/main/docs/api/training.rst

The `step` function provides a lower-level interface for users who require more direct control over the training process in FlowJAX. It allows for customization of individual training steps.

```python
from flowjax.train import step

# Example usage (for advanced control over training steps)
# step(flow, data, ...)
```

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

### Map Posterior and Samples Back to Bounded Domain

Source: https://github.com/danielward27/flowjax/blob/main/docs/examples/snpe.ipynb

After training, this code snippet transforms the learned proposal distribution (now representing the posterior) and the collected parameter samples back to the original bounded domain. This is achieved by applying the inverse of the `to_unbounded` transformation.

```Python
# Undo preprocessing (map everything back to bounded domain)
posterior = Transformed(proposal, bij.Invert(to_unbounded))
data["theta"] = [vmap(to_unbounded.inverse)(t) for t in data["theta"]]
```

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

### Generate Toy Data for Constrained Support

Source: https://github.com/danielward27/flowjax/blob/main/docs/examples/constrained.ipynb

Generates synthetic 2D data where the first dimension (x1) follows a Beta distribution on [0, 1] and the second dimension (x2) is dependent on x1 and normally distributed. This data is used to test constrained flow fitting.

```Python
x1_key, x2_key, key = jr.split(jr.key(0), 3)
x1 = jr.beta(x1_key, a=0.3, b=0.3, shape=(1000,))
x2 = (x1-0.5)**3 + 0.01*jr.normal(x2_key, shape=x1.shape)
x = jnp.stack([x1, x2], axis=-1)
```

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

### Multi-Round SNPE Training Loop

Source: https://github.com/danielward27/flowjax/blob/main/docs/examples/snpe.ipynb

Implements a 3-round Sequential Neural Posterior Estimation (SNPE) training loop. The first round uses MaximumLikelihoodLoss, while subsequent rounds use ContrastiveLoss for potentially more stable posterior estimation. Data is simulated in each round, and the proposal distribution is updated using the fit_to_data function.

```Python
sim_per_round = 1000

data = {"theta": [], "simulations": []}
losses = []
for r in range(3):
    # Carry out simulations (from prior for round 1, the running proposal otherwise).
    key, subkey = jr.split(key)
    if r == 0:
        theta_r = unbounded_prior.sample(subkey, (sim_per_round,))
        loss_fn = MaximumLikelihoodLoss() 
    else:
        theta_r = eqx.filter_jit(proposal.sample)(
            subkey, (sim_per_round,), condition=observed,
        )
        loss_fn = ContrastiveLoss(unbounded_prior, n_contrastive=10)

    key, subkey = jr.split(key)
    simulations_r = task.simulator(
        key=subkey,
        theta=vmap(to_unbounded.inverse)(theta_r),  # Map to bounded for simulating
    )

    data["theta"].append(theta_r)
    data["simulations"].append(simulations_r)

    key, subkey = jr.split(key)
    proposal, losses_r = fit_to_data(
        key=subkey,
        dist=proposal,
        data=(
            jnp.concatenate(data["theta"]),
            jnp.concatenate(data["simulations"]),
        ),
        loss_fn=loss_fn,
        learning_rate=5e-5,
        max_patience=10,
    )
    losses.append(losses_r)
```

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

### Fit Flow to Function using Variational Inference with FlowJAX

Source: https://github.com/danielward27/flowjax/blob/main/docs/api/training.rst

The `fit_to_key_based_loss` function enables fitting a flow model to a specified function using variational inference. This method is suitable when you need to optimize the flow based on a loss function derived from a key.

```python
from flowjax.train import fit_to_key_based_loss

# Example usage (assuming you have a function and a key)
# fit_to_key_based_loss(flow, function, key, ...)
```

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

### Train Flow Without Constraints

Source: https://github.com/danielward27/flowjax/blob/main/docs/examples/constrained.ipynb

Initializes a Masked Autoregressive Flow (MAF) and trains it on the generated toy data without applying any support constraints. This serves as a baseline for comparison.

```Python
key, subkey = jr.split(key)

# Create the flow
unbounded_flow = masked_autoregressive_flow(
    key=subkey,
    base_dist=Normal(jnp.zeros(x.shape[1])))

# Fit the flow, without applying any constraints
key, subkey = jr.split(key)
flow_no_constraints, losses = fit_to_data(
    key=subkey,
    dist=unbounded_flow,
    data=x,
    learning_rate=5e-3,
    max_patience=20,
    max_epochs=70,
)
```

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

### Visualize Posterior Samples with Matplotlib

Source: https://github.com/danielward27/flowjax/blob/main/docs/examples/snpe.ipynb

Generates a series of plots comparing approximated posterior samples against a true posterior. It uses Matplotlib to create scatter plots for each round, customizing titles, aspect ratios, and legends for clarity.

```python
key, subkey = jr.split(key)
reference_posterior = task.sample_reference_posterior(subkey, observed, sim_per_round)

fig, axes = plt.subplots(ncols=len(data["theta"]))

for r, (ax, samps) in enumerate(zip(axes, data["theta"], strict=True)):
    ax.scatter(samps[:, 0], samps[:, 1], s=1, label="Approximation")
    ax.scatter(
        x=reference_posterior[:, 0],
        y=reference_posterior[:, 1],
        s=1,
        label="True posterior",
    )
    title = "Round 0 (prior)" if r == 0 else f"Round {r}"
    ax.set_title(title)
    ax.set_box_aspect(1)
    ax.xaxis.set_ticks([])
    ax.yaxis.set_ticks([])

lgnd = plt.legend(bbox_to_anchor=(1.02, 1), loc="upper left", borderaxespad=0)
for handle in lgnd.legend_handles:
    handle.set_sizes([6.0])

plt.tight_layout()
plt.show()
```

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

### Define Transformation to Unbounded Domain

Source: https://github.com/danielward27/flowjax/blob/main/docs/examples/snpe.ipynb

Creates a transformation pipeline to map data from a bounded domain ([-10, 10]) to an unbounded domain. This involves scaling to [-1, 1] using an Affine transformation and then applying the inverse Tanh function. This is often done to improve numerical stability during model training.

```Python
to_unbounded = bij.Chain(
    [
        bij.Affine(jnp.zeros(task.dim), 1 / task.prior_bound),  # to [-1, 1]
        bij.Invert(bij.Tanh((task.dim,))),  # Arctanh (to unbounded)
    ],
)

unbounded_prior = Transformed(task.prior, to_unbounded)
```

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

### Freeze Parameters in FlowJAX

Source: https://github.com/danielward27/flowjax/blob/main/docs/faq.rst

Demonstrates how to freeze specific parameters in FlowJAX distributions using `paramax.non_trainable`. This utility wraps inexact array leaves with `NonTrainable`, applying `stop_gradient` during unwrapping. It's often used in conjunction with `equinox.tree_at` or `jax.tree_map` to mark parts of a parameter tree as frozen.

```Python
from flowjax.distributions import Normal
import paramax

dist = Normal()
dist = paramax.non_trainable(dist)
```

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

### Train FlowJax Model

Source: https://github.com/danielward27/flowjax/blob/main/docs/examples/unconditional.ipynb

Trains the configured FlowJax model using the generated dataset. The `fit_to_data` function is used with a specified learning rate to optimize the flow's parameters, minimizing the difference between the model's output and the true data distribution.

```Python
key, subkey = jr.split(key)
flow, losses = fit_to_data(subkey, flow, x, learning_rate=1e-3)
```

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

### Create and Train FlowJax Model

Source: https://github.com/danielward27/flowjax/blob/main/docs/examples/conditional.ipynb

Initializes and trains a block neural autoregressive flow model for conditional density estimation. The flow is configured with a Normal base distribution and the specified conditional dimensions. Training is performed using the `fit_to_data` function with specified learning rate and patience.

```Python
key, subkey = jr.split(jr.key(0))
flow = block_neural_autoregressive_flow(
    key=subkey,
    base_dist=Normal(jnp.zeros(x.shape[1])),
    cond_dim=u.shape[1],
)

key, subkey = jr.split(key)
flow, losses = fit_to_data(
    subkey,
    flow,
    data=(x, u),
    learning_rate=5e-2,
    max_patience=10,
)
```

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

### Visualize Learned Posterior Density

Source: https://github.com/danielward27/flowjax/blob/main/docs/examples/variational_inference.ipynb

Visualizes the learned posterior distribution using contour plots and compares it with the true posterior. It includes a helper function 'plot_density' to generate these plots.

```Python
import matplotlib.pyplot as plt
from jax.scipy.stats import multivariate_normal

def plot_density(
    ax, density_fn, xmin=-5, xmax=5, ymin=-5, ymax=5, n=100, levels=None, cmap="Blues",
):
    xvalues = jnp.linspace(xmin, xmax, n)
    yvalues = jnp.linspace(ymin, ymax, n)
    x, y = jnp.meshgrid(xvalues, yvalues)

    points = jnp.hstack([x.reshape(-1, 1), y.reshape(-1, 1)])

    log_prob = density_fn(points).reshape(n, n)
    prob = jnp.exp(log_prob)

    ax.contour(
        prob, levels=levels, extent=[xmin, xmax, ymin, ymax], origin="lower", cmap=cmap,
    )

    ax.set_xlim(xmin, xmax)
    ax.set_ylim(ymin, ymax)


fig, axes = plt.subplots(ncols=2, figsize=(9, 3))
axes[0].set_title("Density plot")

kwargs = {"xmin": 0.25, "xmax": 1.25, "ymin": -1, "ymax": 0, "levels": 5}
plot_density(axes[0], flow.log_prob, cmap="Blues", **kwargs)

# True posterior for comparison
_x = jnp.vstack([jnp.ones_like(x), x])  # full design matrix
cov = jnp.linalg.inv(_x.dot(_x.T) + jnp.eye(2))
mean = cov.dot(_x).dot(y)


def true_posterior_log_prob(theta):
    return multivariate_normal.logpdf(theta, mean, cov)


plot_density(axes[0], true_posterior_log_prob, cmap="Reds", **kwargs)
axes[1].set_title("losses")
axes[1].plot(losses)
```

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

### Append Posterior Samples in Flowjax

Source: https://github.com/danielward27/flowjax/blob/main/docs/examples/snpe.ipynb

Appends posterior samples to the 'theta' key in the data dictionary. It uses JAX for random number generation and sampling from a posterior distribution, conditioned on observed data.

```python
key, subkey = jr.split(key)
data["theta"].append(
    posterior.sample(
        subkey, (sim_per_round,), condition=observed,
    ),
)
```

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

### Generate Toy Data for Conditional Density Estimation

Source: https://github.com/danielward27/flowjax/blob/main/docs/examples/conditional.ipynb

Generates synthetic 2D data for conditional density estimation. It creates a conditioning variable 'u' and a target variable 'x' where the distribution of 'x' depends on 'u', simulating a scenario where the upper limit of a uniform distribution is determined by another uniform random variable.

```Python
key, x_key, cond_key = jr.split(jr.key(0), 3)
u = jr.uniform(cond_key, (10000, 2), minval=0, maxval=5)
x = jr.uniform(x_key, shape=u.shape, maxval=u)
```

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

### Define Bijection for Support Transformation

Source: https://github.com/danielward27/flowjax/blob/main/docs/examples/constrained.ipynb

Defines a bijection using FlowJax's bij module to map data between unbounded and constrained spaces. It uses a Sigmoid transformation for the first dimension to constrain it to [0, 1] and an Identity transformation for the second dimension.

```Python
to_constrained = bij.Stack([
    bij.Sigmoid(),  # x_1: maps real -> [0, 1]
    bij.Identity(), # x_2: maps real -> real
    ])
```

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

### Visualize Learned Conditional Density

Source: https://github.com/danielward27/flowjax/blob/main/docs/examples/conditional.ipynb

Visualizes the learned conditional density of the trained FlowJax model. It samples points on a grid, calculates the probability density for a specific conditioning vector `u = [1, 3]`, and plots the results using a contour plot to show the estimated distribution.

```Python
resolution = 200
test_u = jnp.array([1.0, 3])

xgrid, ygrid = jnp.meshgrid(
    jnp.linspace(-1, 4, resolution), jnp.linspace(-1, 4, resolution),
)
xyinput = jnp.column_stack((xgrid.reshape(-1, 1), ygrid.reshape(-1, 1)))
zgrid = jnp.exp(flow.log_prob(xyinput, test_u).reshape(resolution, resolution))
plt.contourf(xgrid, ygrid, zgrid, levels=50)
plt.show()
```

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

### Generate and Standardize Two-Moons Dataset

Source: https://github.com/danielward27/flowjax/blob/main/docs/examples/unconditional.ipynb

Generates a synthetic dataset approximating the 'two-moons' distribution using FlowJax's `two_moons` function. The generated data is then standardized by subtracting the mean and dividing by the standard deviation for better model performance.

```Python
n_samples = 10000
rng = jr.key(0)
x = two_moons(rng, n_samples)
x = (x - x.mean(axis=0)) / x.std(axis=0)  # Standardize
```

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

### Visualize Regression Fits

Source: https://github.com/danielward27/flowjax/blob/main/docs/examples/variational_inference.ipynb

Visualizes the regression fits generated by sampling from the prior and the trained FlowJax model. It plots multiple regression lines for each model to show the uncertainty in the parameter estimates.

```Python
import matplotlib.pyplot as plt

x_inspect = jnp.linspace(2, -2, n)
plots = [
    ("prior", StandardNormal((2,)), "tab:green"),
    ("trained", flow, "tab:orange"),
]
n_samples = 25

for label, flow, colour in plots:
    key, sample_key = jr.split(key)
    w = flow.sample(sample_key, (n_samples,))
    for ix, (w_0, w_1) in enumerate(w):
        y_inspect = w_0 + w_1 * x_inspect
        lab = label if ix == 0 else None
        plt.plot(x_inspect, y_inspect, alpha=0.3, c=colour, label=lab)

plt.scatter(x, y, label="samples")
plt.title("Sample Fits")
plt.legend()
plt.show()
```

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

### Generate Data for Bayesian Linear Regression

Source: https://github.com/danielward27/flowjax/blob/main/docs/examples/variational_inference.ipynb

Generates synthetic data for a 1D Bayesian linear regression problem. It defines the true parameters (w_0, w_1) and creates input features 'x' and corresponding output 'y' with added noise.

```Python
import jax.numpy as jnp
import jax.random as jr
from jax.scipy.stats import norm

from flowjax.distributions import StandardNormal
from flowjax.flows import masked_autoregressive_flow
from flowjax.train import fit_to_key_based_loss
from flowjax.train.losses import ElboLoss

# generate observed data
data_key = jr.key(0)
w_0 = 0.5
w_1 = -0.5

n = 20
key, x_key, noise_key = jr.split(data_key, 3)
x = jr.uniform(x_key, shape=(n,)) * 4 - 2
y = w_0 + w_1 * x + jr.normal(noise_key, shape=(n,))
```

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

### Train Flow with Constrained Support (Option 1)

Source: https://github.com/danielward27/flowjax/blob/main/docs/examples/constrained.ipynb

Trains a FlowJax model with support constrained to match the target data's support. This involves transforming the flow using a predefined bijection (Sigmoid for the first dimension) and fitting it to slightly clipped data for stability.

```Python
# Contraint option 1: Transform the flow to match data support
flow_constrained_1 = Transformed(
    unbounded_flow,
    non_trainable(to_constrained) # Ensure constraint not trained - doesn't matter here!
)

# Clip for stability
x_clipped = jnp.stack([jnp.clip(x1, 1e-7, 1-1e-7), x2], axis=-1)
flow_constrained_1, losses = fit_to_data(
    key=subkey,
    dist=flow_constrained_1,
    data=x_clipped,
    learning_rate=5e-3,
    max_patience=10,
    max_epochs=70,
)
```

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

### FlowJAX Invert Bijection for Density Evaluation

Source: https://github.com/danielward27/flowjax/blob/main/docs/api/flows.rst

Explains the use of flowjax.bijections.Invert to invert a bijection, prioritizing efficiency in the forward transformation for sampling. This is used by default when transforming the base distribution for density evaluation.

```Python
from flowjax.bijections import Invert

# Example usage (conceptual):
# inverted_bijection = Invert(bijection)
```

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

### Define Proposal Distribution with Masked Autoregressive Flow

Source: https://github.com/danielward27/flowjax/blob/main/docs/examples/snpe.ipynb

Defines the proposal distribution using a masked autoregressive flow (MAF). This flow is configured with a base Normal distribution and an Affine transformer, conditioned on the observed data. MAF is a type of normalizing flow suitable for density estimation.

```Python
key, subkey = jr.split(key)
proposal = masked_autoregressive_flow(
    subkey,
    base_dist=Normal(jnp.zeros(2)),
    transformer=bij.Affine(),
    cond_dim=observed.size,
)
```