boilerplate for pushforwards distributions example

boilerplate.py

import functools

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
import numpyro.distributions as dist
import pandas as pd
import seaborn as sns
from matplotlib import patches

sns.set()

mu = jnp.zeros(2)
cov = jnp.array([[0.5, 0.2], [0.2, 0.1]])
distribution = dist.MultivariateNormal(mu, cov)
rng_key = jax.random.PRNGKey(2)


def g(x):
    return jnp.asarray([jax.scipy.special.expit(x[0]), jnp.exp(x[1])])